From 910465451562243ca9f23a430d1a20c9b1753e4e Mon Sep 17 00:00:00 2001
From: Faisal Anees <faisal.iiit@gmail.com>
Date: Mon, 8 Feb 2021 10:14:39 +0400
Subject: [PATCH] Read MLFLow creds securely

---
 hydra/docker/Dockerfile |  1 +
 hydra/docker/entry.py   | 15 ++++++++++++++-
 hydra/utils/secrets.py  | 25 +++++++++++++++++++++++++
 requirements.txt        |  3 ++-
 4 files changed, 42 insertions(+), 2 deletions(-)
 create mode 100644 hydra/utils/secrets.py

diff --git a/hydra/docker/Dockerfile b/hydra/docker/Dockerfile
index 5711d8d..6c0977e 100644
--- a/hydra/docker/Dockerfile
+++ b/hydra/docker/Dockerfile
@@ -1,4 +1,5 @@
 FROM continuumio/miniconda3
 WORKDIR /home
 COPY entry.py .
+RUN pip install hydra-ml==0.3.8
 ENTRYPOINT ["python", "entry.py"]
diff --git a/hydra/docker/entry.py b/hydra/docker/entry.py
index 2a392d1..760d757 100644
--- a/hydra/docker/entry.py
+++ b/hydra/docker/entry.py
@@ -3,6 +3,8 @@
 import argparse
 import subprocess
 
+from hydra.utils.secrets import get_creds_for_gcp_mlflow
+
 CONDA_ENV_NAME = "hydra"
 
 args_parser = argparse.ArgumentParser()
@@ -30,13 +32,24 @@
     shutil.copytree("/home/data", "/home/project/data")
 
 subprocess.run(["conda", "env", "create", "-n", CONDA_ENV_NAME, "-f", "environment.yml"])
-subprocess.run(["conda", "run", "-n", "hydra", "pip", "install", "hydra-ml"])
+subprocess.run(["conda", "run", "-n", "hydra", "pip", "install", "hydra-ml==0.3.8"])
 
 if args.options is not None:
     for arg in args.options.split():
         [key, val] = arg.split('=')
         os.putenv(key, val)
 
+mlflow_tracking_uri, mlflow_username,\
+    mlflow_pswd = "", "", ""
+
+if os.environ.get('HYDRA_PLATFORM') == 'gcp':
+    mlflow_tracking_uri, mlflow_username,\
+    mlflow_pswd = get_creds_for_gcp_mlflow()
+
+os.putenv('MLFLOW_TRACKING_URI', mlflow_tracking_uri)
+os.putenv('MLFLOW_USERNAME', mlflow_username)
+os.putenv('MLFLOW_PASSWORD', mlflow_pswd)
+
 os.putenv('HYDRA_PLATFORM', args.platform)
 os.putenv('HYDRA_GIT_URL', args.git_url)
 os.putenv('HYDRA_COMMIT_SHA', args.commit_sha)
diff --git a/hydra/utils/secrets.py b/hydra/utils/secrets.py
new file mode 100644
index 0000000..f33acb8
--- /dev/null
+++ b/hydra/utils/secrets.py
@@ -0,0 +1,25 @@
+import os
+
+from google.cloud import secretmanager
+
+def access_secret_version(project_id, secret_id, version_id="latest"):
+    # Create the Secret Manager client.
+    client = secretmanager.SecretManagerServiceClient()
+
+    # Build the resource name of the secret version.
+    name = f"projects/{project_id}/secrets/{secret_id}/versions/{version_id}"
+
+    # Access the secret version.
+    response = client.access_secret_version(name=name)
+
+    # Return the decoded payload.
+    return response.payload.data.decode('UTF-8')
+
+
+def get_creds_for_gcp_mlflow():
+    project_id = os.environ["GCP_PROJECT"]
+
+    tracking_uri = access_secret_version(project_id, 'MLFLOW_TRACKING_URI')
+    username = access_secret_version(project_id, 'MLFLOW_TRACKING_USERNAME')
+    pswd = access_secret_version(project_id, 'MLFLOW_TRACKING_PASSWORD')
+    return tracking_uri, username, pswd
diff --git a/requirements.txt b/requirements.txt
index fa4ea9d..02f4ee2 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -6,4 +6,5 @@ GitPython==3.1.9
 google-cloud-storage==1.32.0
 docker==4.3.1
 pyyaml
-boto3==1.16.28
\ No newline at end of file
+boto3==1.16.28
+google-cloud-secret-manager==2.2.0
\ No newline at end of file