Repository: tensorflow/tfx

Issue:
Kubeflow Runner/BaseComponent dont force inject GCP credentials/secret into pods
# Description
As many people we don’t deploy Kubeflow on GCP, therefore we don’t need/have GCP cred. Unfortunately the GCP env vars and more annoyingly the secret is hardcoded into the BaseComponent. Therefore no component pod can be started without the required secret (named _user-gcp-sa_ per default). 

# Possible Solution
Please make this mount optionally, controlled via an option given to the runner. 
This can go in hand with a better documentation of the additional arguments of the Pipeline class. 
A possible solution is to create an overall platform parameter, similar to the solution in the [Pipelines example](https://github.com/kubeflow/pipelines/blob/master/samples/tfx/taxi-cab-classification-pipeline.py): 

The relevant code is around line 124 in` tfx.orchestration.kubeflow.base_component`.  I am happy to assist if needed. 

# Workaround
Deploy a fake secret:
```
apiVersion: v1
kind: Secret
metadata:
  name: user-gcp-sa
type: Opaque
data:
  user-gcp-sa.json: YXBpVXJsOiAiaHR0cHM6Ly9teS5hcGkuY29tL2FwaS92MSIKdXNlcm5hbWU6IHt7dXNlcm5hbWV9fQpwYXNzd29yZDoge3twYXNzd29yZH19

```


Discussion:
We'll make this optional. I'll send out a fix shortly.

Resolving patch:
diff --git a/tfx/dependencies.py b/tfx/dependencies.py
--- a/tfx/dependencies.py
+++ b/tfx/dependencies.py
@@ -36,7 +36,7 @@ def make_required_test_packages():
   return [
       'apache-airflow>=1.10,<2',
       'docker>=3.7,<4',
-      'kfp>=0.1,<=0.1.11; python_version >= "3.0"',
+      'kfp>=0.1.18,<0.2; python_version >= "3.0"',
       'pytest>=4.4.1,<5',
       'tensorflow>=1.13,<2',
       'tzlocal>=1.5,<2.0',
diff --git a/tfx/orchestration/kubeflow/base_component.py b/tfx/orchestration/kubeflow/base_component.py
--- a/tfx/orchestration/kubeflow/base_component.py
+++ b/tfx/orchestration/kubeflow/base_component.py
@@ -28,7 +28,6 @@
 import collections
 import json
 from kfp import dsl
-from kfp import gcp
 from kubernetes import client as k8s_client
 from typing import Any, Dict, List, Optional, Text
 
@@ -123,7 +122,7 @@ def __new__(
         image=pipeline_properties.tfx_image,
         arguments=arguments,
         file_outputs=file_outputs,
-    ).apply(gcp.use_gcp_secret('user-gcp-sa'))  # Adds GCP authentication.
+    )
 
     # Add the Argo workflow ID to the container's environment variable so it
     # can be used to uniquely place pipeline outputs under the pipeline_root.
diff --git a/tfx/orchestration/kubeflow/runner.py b/tfx/orchestration/kubeflow/runner.py
--- a/tfx/orchestration/kubeflow/runner.py
+++ b/tfx/orchestration/kubeflow/runner.py
@@ -21,13 +21,70 @@
 
 from kfp import compiler
 from kfp import dsl
-from typing import Optional, Text
+from kfp import gcp
+from typing import Callable, List, Optional, Text
 
 from tfx.components.base import base_component as tfx_base_component
 from tfx.orchestration import pipeline as tfx_pipeline
 from tfx.orchestration import tfx_runner
 from tfx.orchestration.kubeflow import base_component
 
+# OpFunc represents the type of a function that takes as input a
+# dsl.ContainerOp and returns the same object. Common operations such as adding
+# k8s secrets, mounting volumes, specifying the use of TPUs and so on can be
+# specified as an OpFunc.
+# See example usage here:
+# https://github.com/kubeflow/pipelines/blob/master/sdk/python/kfp/gcp.py
+OpFunc = Callable[[dsl.ContainerOp], dsl.ContainerOp]
+
+
+# Default secret name for GCP credentials. This secret is installed as part of
+# a typical Kubeflow installation when the platform is GKE.
+_KUBEFLOW_GCP_SECRET_NAME = 'user-gcp-sa'
+
+
+def get_default_pipeline_operator_funcs() -> List[OpFunc]:
+  """Returns a default list of pipeline operator functions.
+
+  Returns:
+    A list of functions with type OpFunc.
+  """
+  # Enables authentication for GCP services in a typical GKE Kubeflow
+  # installation.
+  gcp_secret_op = gcp.use_gcp_secret(_KUBEFLOW_GCP_SECRET_NAME)
+
+  return [gcp_secret_op]
+
+
+class KubeflowRunnerConfig(object):
+  """Runtime configuration parameters specific to execution on Kubeflow."""
+
+  def __init__(self, pipeline_operator_funcs: Optional[List[OpFunc]] = None):
+    """Creates a KubeflowRunnerConfig object.
+
+    The user can use pipeline_operator_funcs to apply modifications to
+    ContainerOps used in the pipeline. For example, to ensure the pipeline
+    steps mount a GCP secret, and a Persistent Volume, one can create config
+    object like so:
+
+      from kfp import gcp, onprem
+      mount_secret_op = gcp.use_secret('my-secret-name)
+      mount_volume_op = onprem.mount_pvc(
+        "my-persistent-volume-claim",
+        "my-volume-name",
+        "/mnt/volume-mount-path")
+
+      config = KubeflowRunnerConfig(
+        pipeline_operator_funcs=[mount_secret_op, mount_volume_op]
+      )
+
+    Args:
+      pipeline_operator_funcs: A list of ContainerOp modifying functions that
+        will be applied to every container step in the pipeline.
+    """
+    self.pipeline_operator_funcs = (
+        pipeline_operator_funcs or get_default_pipeline_operator_funcs())
+
 
 class KubeflowRunner(tfx_runner.TfxRunner):
   """Kubeflow Pipelines runner.
@@ -35,14 +92,19 @@ class KubeflowRunner(tfx_runner.TfxRunner):
   Constructs a pipeline definition YAML file based on the TFX logical pipeline.
   """
 
-  def __init__(self, output_dir: Optional[Text] = None):
+  def __init__(self,
+               output_dir: Optional[Text] = None,
+               config: Optional[KubeflowRunnerConfig] = None):
     """Initializes KubeflowRunner for compiling a Kubeflow Pipeline.
 
     Args:
       output_dir: An optional output directory into which to output the pipeline
         definition files. Defaults to the current working directory.
+      config: An optional KubeflowRunnerConfig object to specify runtime
+        configuration when running the pipeline under Kubeflow.
     """
     self._output_dir = output_dir or os.getcwd()
+    self._config = config or KubeflowRunnerConfig()
 
   def _prepare_output_dict(self, outputs: tfx_base_component.ComponentOutputs):
     return dict((k, v.get()) for k, v in outputs.get_all().items())
@@ -104,6 +166,9 @@ def _construct_pipeline_graph(self, pipeline: tfx_pipeline.Pipeline):
           executor_class_path=executor_class_path,
           pipeline_properties=pipeline_properties)
 
+      for operator in self._config.pipeline_operator_funcs:
+        kfp_component.container_op.apply(operator)
+
       for channel_name, channel in component.outputs.get_all().items():
         producers[channel] = {}
         producers[channel]['component'] = kfp_component
@@ -129,5 +194,8 @@ def _construct_pipeline():
       self._construct_pipeline_graph(pipeline)
 
     pipeline_name = pipeline.pipeline_args['pipeline_name']
+    # TODO(b/134680219): Allow users to specify the extension. Specifying
+    # .yaml will compile the pipeline directly into a YAML file. Kubeflow
+    # backend recognizes .tar.gz, .zip, and .yaml today.
     pipeline_file = os.path.join(self._output_dir, pipeline_name + '.tar.gz')
     compiler.Compiler().compile(_construct_pipeline, pipeline_file)


Test patch:
diff --git a/tfx/orchestration/kubeflow/runner_test.py b/tfx/orchestration/kubeflow/runner_test.py
--- a/tfx/orchestration/kubeflow/runner_test.py
+++ b/tfx/orchestration/kubeflow/runner_test.py
@@ -21,6 +21,7 @@
 import shutil
 import tarfile
 import tempfile
+from kfp import onprem
 import tensorflow as tf
 import yaml
 
@@ -57,8 +58,7 @@ def test_two_step_pipeline(self):
     """Sanity-checks the construction and dependencies for a 2-step pipeline.
     """
     runner.KubeflowRunner().run(_two_step_pipeline())
-    file_path = os.path.join(self.test_dir,
-                             'two_step_pipeline.tar.gz')
+    file_path = os.path.join(self.test_dir, 'two_step_pipeline.tar.gz')
     self.assertTrue(tf.gfile.Exists(file_path))
 
     with tarfile.TarFile.open(file_path).extractfile(
@@ -72,9 +72,7 @@ def test_two_step_pipeline(self):
       self.assertEqual(2, len(containers))
 
       big_query_container = [
-          c
-          for c in containers
-          if c['name'] == 'bigqueryexamplegen'
+          c for c in containers if c['name'] == 'bigqueryexamplegen'
       ]
       self.assertEqual(1, len(big_query_container))
       self.assertEqual([
@@ -83,9 +81,7 @@ def test_two_step_pipeline(self):
       ], big_query_container[0]['container']['command'])
 
       statistics_gen_container = [
-          c
-          for c in containers
-          if c['name'] == 'statisticsgen'
+          c for c in containers if c['name'] == 'statisticsgen'
       ]
       self.assertEqual(1, len(statistics_gen_container))
 
@@ -100,20 +96,94 @@ def test_two_step_pipeline(self):
           'tasks': [{
               'name': 'bigqueryexamplegen',
               'template': 'bigqueryexamplegen'
-          },
-                    {
-                        'name': 'statisticsgen',
-                        'template': 'statisticsgen',
-                        'dependencies': ['bigqueryexamplegen'],
-                        'arguments': {
-                            'parameters': [{
-                                'name': 'bigqueryexamplegen-examples',
-                                'value': parameter_value
-                            }],
-                        },
-                    }]
+          }, {
+              'name': 'statisticsgen',
+              'template': 'statisticsgen',
+              'dependencies': ['bigqueryexamplegen'],
+              'arguments': {
+                  'parameters': [{
+                      'name': 'bigqueryexamplegen-examples',
+                      'value': parameter_value
+                  }],
+              },
+          }]
       }, dag[0]['dag'])
 
+  def test_default_pipeline_operator_funcs(self):
+    runner.KubeflowRunner().run(_two_step_pipeline())
+    file_path = os.path.join(self.test_dir, 'two_step_pipeline.tar.gz')
+    self.assertTrue(tf.gfile.Exists(file_path))
+
+    with tarfile.TarFile.open(file_path).extractfile(
+        'pipeline.yaml') as pipeline_file:
+      self.assertIsNotNone(pipeline_file)
+      pipeline = yaml.load(pipeline_file)
+
+      containers = [
+          c for c in pipeline['spec']['templates'] if 'container' in c
+      ]
+      self.assertEqual(2, len(containers))
+
+      # Check that each container has default GCP credentials.
+
+      container_0 = containers[0]
+      env = [
+          env for env in container_0['container']['env']
+          if env['name'] == 'GOOGLE_APPLICATION_CREDENTIALS'
+      ]
+      self.assertEqual(1, len(env))
+      self.assertEqual('/secret/gcp-credentials/user-gcp-sa.json',
+                       env[0]['value'])
+
+      container_1 = containers[0]
+      env = [
+          env for env in container_1['container']['env']
+          if env['name'] == 'GOOGLE_APPLICATION_CREDENTIALS'
+      ]
+      self.assertEqual(1, len(env))
+      self.assertEqual('/secret/gcp-credentials/user-gcp-sa.json',
+                       env[0]['value'])
+
+  def test_volume_mounting_pipeline_operator_funcs(self):
+    mount_volume_op = onprem.mount_pvc('my-persistent-volume-claim',
+                                       'my-volume-name',
+                                       '/mnt/volume-mount-path')
+    config = runner.KubeflowRunnerConfig(
+        pipeline_operator_funcs=[mount_volume_op])
+
+    runner.KubeflowRunner(config=config).run(_two_step_pipeline())
+    file_path = os.path.join(self.test_dir, 'two_step_pipeline.tar.gz')
+    self.assertTrue(tf.gfile.Exists(file_path))
+
+    with tarfile.TarFile.open(file_path).extractfile(
+        'pipeline.yaml') as pipeline_file:
+      self.assertIsNotNone(pipeline_file)
+      pipeline = yaml.load(pipeline_file)
+
+      containers = [
+          c for c in pipeline['spec']['templates'] if 'container' in c
+      ]
+      self.assertEqual(2, len(containers))
+
+      # Check that each container has the volume mounted.
+      self.assertEqual([{
+          'name': 'my-volume-name',
+          'mountPath': '/mnt/volume-mount-path'
+      }], containers[0]['container']['volumeMounts'])
+
+      self.assertEqual([{
+          'name': 'my-volume-name',
+          'mountPath': '/mnt/volume-mount-path'
+      }], containers[1]['container']['volumeMounts'])
+
+      # Check that the PVC is specified.
+      self.assertEqual([{
+          'name': 'my-volume-name',
+          'persistentVolumeClaim': {
+              'claimName': 'my-persistent-volume-claim'
+          }
+      }], pipeline['spec']['volumes'])
+
 
 if __name__ == '__main__':
   tf.test.main()