add_pod_env op handler

sdk/python/kfp/compiler/_default_transformers.py

@@ -0,0 +1,41 @@
+# 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.
+
+from ..dsl._container_op import BaseOp
+
+def add_pod_env(op: BaseOp) -> BaseOp:
+    """Adds pod environment info to the operator.
+    """
+    if op.pod_labels and op.pod_labels['add-pod-env'] == 'true':
+        from kubernetes import client as k8s_client
+        op.container.add_env_variable(
+            k8s_client.V1EnvVar(
+                name='KFP_POD_NAME', 
+                value_from=k8s_client.V1EnvVarSource(
+                    field_ref=k8s_client.V1ObjectFieldSelector(
+                        field_path='metadata.name'
+                    )
+                )
+            )
+        ).add_env_variable(
+            k8s_client.V1EnvVar(
+                name='KFP_NAMESPACE', 
+                value_from=k8s_client.V1EnvVarSource(
+                    field_ref=k8s_client.V1ObjectFieldSelector(
+                        field_path='metadata.namespace'
+                    )
+                )
+            )
+        )
+    return op

sdk/python/kfp/compiler/compiler.py

@@ -23,6 +23,7 @@
 from .. import dsl
 from ._k8s_helper import K8sHelper
 from ._op_to_template import _op_to_template
+from ._default_transformers import add_pod_env
 
 from ..dsl._metadata import TypeMeta, _extract_pipeline_metadata
 from ..dsl._ops_group import OpsGroup
@@ -652,7 +653,9 @@ def _compile(self, pipeline_func):
       sanitized_ops[sanitized_name] = op
     p.ops = sanitized_ops
 
-    workflow = self._create_pipeline_workflow(args_list_with_defaults, p, p.conf.op_transformers)
+    op_transformers = [add_pod_env]
+    op_transformers.extend(p.conf.op_transformers)
+    workflow = self._create_pipeline_workflow(args_list_with_defaults, p, op_transformers)
     return workflow
 
   def compile(self, pipeline_func, package_path, type_check=True):

sdk/python/tests/compiler/compiler_tests.py

@@ -548,3 +548,6 @@ def some_pipeline():
       container = template.get('container', None)
       if container:
         self.assertEqual(template['retryStrategy']['limit'], 5)
+
+  def test_add_pod_env(self):
+    self._test_py_compile_yaml('add_pod_env')

sdk/python/tests/compiler/testdata/add_pod_env.py

@@ -0,0 +1,32 @@
+# 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.
+
+import kfp
+import kfp.dsl as dsl
+
+@dsl.pipeline(
+  name='Test adding pod env',
+  description='Test adding pod env'
+)
+def test_add_pod_env():
+    op = dsl.ContainerOp(
+     name='echo',
+     image='library/bash',
+     command=['sh', '-c'],
+     arguments=['echo $KFP_POD_NAME']).add_pod_label(
+         'add-pod-env', 'true'
+     )
+
+if __name__ == '__main__':
+    kfp.compiler.Compiler().compile(test_add_pod_env, __file__ + '.yaml')

sdk/python/tests/compiler/testdata/add_pod_env.yaml

@@ -0,0 +1,43 @@
+apiVersion: argoproj.io/v1alpha1
+kind: Workflow
+metadata:
+  generateName: test-adding-pod-env-
+spec:
+  arguments:
+    parameters: []
+  entrypoint: test-adding-pod-env
+  serviceAccountName: pipeline-runner
+  templates:
+  - container:
+      args:
+      - echo $KFP_POD_NAME
+      command:
+      - sh
+      - -c
+      env:
+      - name: KFP_POD_NAME
+        valueFrom:
+          fieldRef:
+            fieldPath: metadata.name
+      - name: KFP_NAMESPACE
+        valueFrom:
+          fieldRef:
+            fieldPath: metadata.namespace
+      image: library/bash
+    metadata:
+      labels:
+        add-pod-env: 'true'
+    name: echo
+    outputs:
+      artifacts:
+      - name: mlpipeline-ui-metadata
+        optional: true
+        path: /mlpipeline-ui-metadata.json
+      - name: mlpipeline-metrics
+        optional: true
+        path: /mlpipeline-metrics.json
+  - dag:
+      tasks:
+      - name: echo
+        template: echo
+    name: test-adding-pod-env