diff --git a/sdk/python/kfp/compiler/_default_transformers.py b/sdk/python/kfp/compiler/_default_transformers.py new file mode 100644 index 00000000000..34711be6b58 --- /dev/null +++ b/sdk/python/kfp/compiler/_default_transformers.py @@ -0,0 +1,41 @@ +# Copyright 2019 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, ContainerOp + +def add_pod_env(op: BaseOp) -> BaseOp: + """Adds pod environment info to ContainerOp. + """ + if isinstance(op, ContainerOp) and 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 \ No newline at end of file diff --git a/sdk/python/kfp/compiler/compiler.py b/sdk/python/kfp/compiler/compiler.py index 0a8b2df1509..67855bffb96 100644 --- a/sdk/python/kfp/compiler/compiler.py +++ b/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): diff --git a/sdk/python/tests/compiler/compiler_tests.py b/sdk/python/tests/compiler/compiler_tests.py index dd16ca15025..b8c9da359cc 100644 --- a/sdk/python/tests/compiler/compiler_tests.py +++ b/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') diff --git a/sdk/python/tests/compiler/testdata/add_pod_env.py b/sdk/python/tests/compiler/testdata/add_pod_env.py new file mode 100644 index 00000000000..584498a6a7e --- /dev/null +++ b/sdk/python/tests/compiler/testdata/add_pod_env.py @@ -0,0 +1,32 @@ +# Copyright 2019 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') \ No newline at end of file diff --git a/sdk/python/tests/compiler/testdata/add_pod_env.yaml b/sdk/python/tests/compiler/testdata/add_pod_env.yaml new file mode 100644 index 00000000000..fb4b114587e --- /dev/null +++ b/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