Add support for AWS access/secret keys in train component (#466)

mnist/ks_app/components/params.libsonnet

@@ -12,6 +12,7 @@
       numPs: 0,
       numWorkers: 0,
       secret: '',
+      secretKeyRefs: '',
       trainSteps: 200,
     },
     "mnist-deploy-gcp": {

mnist/ks_app/components/train.jsonnet

@@ -43,6 +43,28 @@ local trainEnv = [
   },
 ];
 
+// AWS Access/Secret keys
+local awsSecretKeyRefs = util.parseSecret(params.secretKeyRefs);
+local awsAccessKeyId = if std.length(awsSecretKeyRefs) > 0 then awsSecretKeyRefs[0] else "";
+local awsSecretAccessKey = if std.length(awsSecretKeyRefs) > 1 then awsSecretKeyRefs[1] else "";
+
+local awsEnv = [
+  {
+    name: "AWS_ACCESS_KEY_ID",
+    valueFrom: {
+      secretKeyRef:  
+        awsAccessKeyId
+    }
+  },
+  {
+    name: "AWS_SECRET_ACCESS_KEY",
+    valueFrom: {
+      secretKeyRef:
+        awsSecretAccessKey
+    }
+  }
+];
+
 local secretPieces = std.split(params.secret, "=");
 local secretName = if std.length(secretPieces) > 0 then secretPieces[0] else "";
 local secretMountPath = if std.length(secretPieces) > 1 then secretPieces[1] else "";
@@ -54,7 +76,7 @@ local replicaSpec = {
         "/usr/bin/python",
         "/opt/model.py",
       ],
-      env: trainEnv + util.parseEnv(params.envVariables),
+      env: trainEnv + util.parseEnv(params.envVariables) + awsEnv,
       image: params.image,
       name: "tensorflow",
       volumeMounts: if secretMountPath != "" then

mnist/ks_app/components/util.libsonnet

@@ -7,6 +7,13 @@
       value: v[1],
     },
 
+  // convert a list of two items into a map representing a secret name and key
+  listToSecretMap:: function(v)
+    {
+      name: v[0],
+      key: v[1],
+    },
+
   // Function to turn comma separated list of environment variables into a dictionary.
   parseEnv:: function(v)
     local pieces = std.split(v, ",");
@@ -16,4 +23,14 @@
         std.split(v, ",")
       )
     else [],
+
+  // Function to turn comma separated list of secret names and keys into a dictionary.
+  parseSecret:: function(v)
+    local pieces = std.split(v, ",");
+    if v != "" && std.length(pieces) > 0 then
+      std.map(
+        function(i) $.listToSecretMap(std.split(i, ".")),
+        std.split(v, ",")
+      )
+    else [],
 }