diff --git a/mnist/ks_app/components/train.jsonnet b/mnist/ks_app/components/train.jsonnet
index d8b76f02..3dd0ac19 100644
--- a/mnist/ks_app/components/train.jsonnet
+++ b/mnist/ks_app/components/train.jsonnet
@@ -44,26 +44,7 @@ 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 trainSecrets = util.parseSecrets(params.secretKeyRefs);
 
 local secretPieces = std.split(params.secret, "=");
 local secretName = if std.length(secretPieces) > 0 then secretPieces[0] else "";
@@ -76,7 +57,7 @@ local replicaSpec = {
         "/usr/bin/python",
         "/opt/model.py",
       ],
-      env: trainEnv + util.parseEnv(params.envVariables) + awsEnv,
+      env: trainEnv + util.parseEnv(params.envVariables) + trainSecrets,
       image: params.image,
       name: "tensorflow",
       volumeMounts: if secretMountPath != "" then
diff --git a/mnist/ks_app/components/util.libsonnet b/mnist/ks_app/components/util.libsonnet
index b6d7a139..82334b92 100644
--- a/mnist/ks_app/components/util.libsonnet
+++ b/mnist/ks_app/components/util.libsonnet
@@ -7,11 +7,16 @@
       value: v[1],
     },
 
-  // convert a list of two items into a map representing a secret name and key
+  // convert a list of two items into a map representing an env variable referencing k8s secret
   listToSecretMap:: function(v)
     {
       name: v[0],
-      key: v[1],
+        valueFrom: {
+          secretKeyRef: {
+            name: std.split(v[1], ".")[0],
+            key: std.split(v[1], ".")[1],
+          }
+        }
     },
 
   // Function to turn comma separated list of environment variables into a dictionary.
@@ -24,12 +29,12 @@
       )
     else [],
 
-  // Function to turn comma separated list of secret names and keys into a dictionary.
-  parseSecret:: function(v)
+  // Function to turn comma separated list of env variables referencing secrets into a dictionary.
+  parseSecrets:: function(v)
     local pieces = std.split(v, ",");
     if v != "" && std.length(pieces) > 0 then
       std.map(
-        function(i) $.listToSecretMap(std.split(i, ".")),
+        function(i) $.listToSecretMap(std.split(i, "=")),
         std.split(v, ",")
       )
     else [],