SDK/GCP - Replaced default_gcp_op with task.apply(use_gcp_secret)

sdk/python/kfp/dsl/__init__.py

@@ -18,4 +18,3 @@
 from ._container_op import ContainerOp
 from ._ops_group import OpsGroup, ExitHandler, Condition
 from ._component import python_component
-from ._default_gcp_op import default_gcp_op

sdk/python/kfp/dsl/_container_op.py

@@ -88,6 +88,21 @@ def __init__(self, name: str, image: str, command: str=None, arguments: str=None
     if len(self.outputs) == 1:
       self.output = list(self.outputs.values())[0]
 
+  def apply(self, mod_func):
+    """Applies a modifier function to self. The function should return the passed object.
+    This is needed to chain "extention methods" to this class.
+
+    Example:
+      from kfp.gcp import use_gcp_secret
+      task = (
+        train_op(...)
+          .set_memory_request('1GB')
+          .apply(use_gcp_secret('user-gcp-sa'))
+          .set_memory_limit('2GB')
+      )
+    """
+    return mod_func(self)
+
   def after(self, op):
     """Specify explicit dependency on another op."""
     self.dependent_op_names.append(op.name)

sdk/python/kfp/gcp.py

@@ -0,0 +1,58 @@
+# Copyright 2018 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+def use_gcp_secret(secret_name='user-gcp-sa', secret_file_path_in_volume='/user-gcp-sa.json', volume_name='gcp-credentials', secret_volume_mount_path='/secret/gcp-credentials'):
+    """An operator that configures the container to use GCP service account.
+
+        The user-gcp-sa secret is created as part of the kubeflow deployment that
+        stores the access token for kubeflow user service account.
+
+        With this service account, the container has a range of GCP APIs to
+        access to. This service account is automatically created as part of the
+        kubeflow deployment.
+
+        For the list of the GCP APIs this service account can access to, check
+        https://github.com/kubeflow/kubeflow/blob/7b0db0d92d65c0746ac52b000cbc290dac7c62b1/deployment/gke/deployment_manager_configs/iam_bindings_template.yaml#L18
+
+        If you want to call the GCP APIs in a different project, grant the kf-user
+        service account access permission.
+    """
+
+    def _use_gcp_secret(task):
+        from kubernetes import client as k8s_client
+        return (
+            task
+                .add_volume(
+                    k8s_client.V1Volume(
+                        name=volume_name,
+                        secret=k8s_client.V1SecretVolumeSource(
+                            secret_name=secret_name,
+                        )
+                    )
+                )
+                .add_volume_mount(
+                    k8s_client.V1VolumeMount(
+                        name=volume_name,
+                        mount_path=secret_volume_mount_path,
+                    )
+                )
+                .add_env_variable(
+                    k8s_client.V1EnvVar(
+                        name='GOOGLE_APPLICATION_CREDENTIALS',
+                        value=secret_volume_mount_path + secret_file_path_in_volume,
+                    )
+                )
+        )
+
+    return _use_gcp_secret