How can I mount a GCS bucket in a custom Docker image on AI Platform?

Question

I'm using Google's AI Platform to train machine learning models using a custom Docker image. To run existing code without modifications, I would like to mount a GCS bucket inside the container.

I think one way to achieve this is to install gcloud to authentication and gcsfuse for mounting in the container. My Dockerfile looks like this:

FROM nvidia/cuda:10.1-cudnn7-runtime-ubuntu18.04

WORKDIR /root

# Install system packages.
RUN apt-get update
RUN apt-get install -y curl
# ...

# Install gcsfuse.
RUN echo "deb http://packages.cloud.google.com/apt gcsfuse-bionic main" | tee /etc/apt/sources.list.d/gcsfuse.list
RUN curl https://packages.cloud.google.com/apt/doc/apt-key.gpg | apt-key add -
RUN apt-get update
RUN apt-get install -y gcsfuse

# Install gcloud.
RUN apt-get install -y apt-transport-https
RUN apt-get install -y ca-certificates
RUN echo "deb [signed-by=/usr/share/keyrings/cloud.google.gpg] https://packages.cloud.google.com/apt cloud-sdk main" | tee -a /etc/apt/sources.list.d/google-cloud-sdk.list
RUN curl https://packages.cloud.google.com/apt/doc/apt-key.gpg | apt-key --keyring /usr/share/keyrings/cloud.google.gpg add -
RUN apt-get update
RUN apt-get install -y google-cloud-sdk

# ...

ENTRYPOINT ["entrypoint.sh"]

Inside the entry point script, I then try to authenticate with Google cloud and mount the bucket. My entrypoint.sh looks like this:

#!/bin/sh
set -e

gcloud auth login
gcsfuse my-bucket-name /root/output
python3 script.py --logdir /root/output/experiment

I then build the container and run it either locally for testing or remotely on the AI Platform for the full training run:

# Run locally for testing.
nvidia-docker build -t my-image-name .
nvidia-docker run -it --rm my-image-name

# Run on AI Platform for full training run.
nvidia-docker build -t my-image-name .
gcloud auth configure-docker
nvidia-docker push my-image-name
gcloud beta ai-platform jobs submit training --region us-west1 --scale-tier custom --master-machine-type standard_p100 --master-image-uri my-image-name

Both locally and on the AI Platform, the entrypoint.sh script hangs at the line gcloud auth login, probably because it waits for user input. Is there a better way of authenticating with Google Cloud from within the container? If not, how can I automate the line that currently hangs?

user_1738915 · Answer

Instead of using gcloud auth login which is primarily meant for human/user authentication, consider using gcloud auth activate-service-account and supplying a key file. See here for details:

https://cloud.google.com/sdk/gcloud/reference/auth/activate-service-account

I would recommend not placing the keys file inside the image but instead provide it externally. Another alternative is to realize that the authentication can implicit via environment variables. So following cloud native practices, have the environment provide the credentials needed and don't try and authenticate inside your environment at all. If you plan to run your container inside GCP Compute Engine or GKE you can implicitly provide the service account to the container from outside the container.

user_1054322 · Answer

If the default service account meets your needs, you can configure your container to use it like this. You may also be able to give it what it needs by granting it extra permissions.

If you want to use your own service account, you'll need to authenticate as a service account via:

gcloud auth activate-service-account --key-file=somekey.json

That way the container won't hang while asking you to authenticate via a browser. So the obvious next question is:

How do I insert my service account's key into the container?

The Strategy

First, you'll want to generate a key file for whatever service account you do want to use.

It's not a good idea to store credentials in docker images, so I put the key in a script which I then put in a storage bucket. So the container downloads and runs the script, which switches the configured identity to a service account of my choosing.

Entrypoint

# runs as the default service account
gsutil cp "$1" /run/cmd
chmod +x /run/cmd
/run/cmd

Run Script (in bucket)

cat << EOF!! > /dev/shm/sa_key
THE KEY FILE CONTENTS GO HERE
EOF!!

gcloud auth activate-service-account --key-file=/dev/shm/sa_key

# commands below this line are performed with the specified identity

The default service account has access to the storage buckets in its project, so the script above will have to go in such a bucket. Be sure that that bucket is appropriately protected, anyone with access to it can assume the identity of the service account whose keys it contains.

Testing Locally

docker run -v "/home/me/.config/gcloud:/root/.config/gcloud" \
    theimagename gs://my-project_job1/run_script

This will use your user's active gcloud creds to pull down the script and then it will switch to the service account. When it finishes, your host's gcloud will be configured to use the service account--so you may need to switch it back to yourself vi gcloud auth login. To avoid this, you can instead mount a copy of that directory, that way the original remains untouched.

Running in GCP

gcloud ai-platform jobs submit training job1 \    
  --region us-west2 \
  --master-image-uri us.gcr.io/my-project/theimagename:latest \
  -- gs://my-project_job1/run_script

I hacked this up a bit to remove references to parts of my project that are irrelevant here, so this probably won't run as is, but I think this shows the gist of how I've been using it:

https://gist.github.com/MatrixManAtYrService/737cb408e5a27c2aaa19576b0f6ec18a