From a77d8e9d03ac09cd938268c61e52d81830a0c2dd Mon Sep 17 00:00:00 2001 From: Ilias Katsakioris Date: Sat, 5 Oct 2019 05:59:12 +0300 Subject: [PATCH] SDK/DSL: ContainerOp.add_pvolume - Fix volume passed in add_volume (#2306) Signed-off-by: Ilias Katsakioris --- sdk/python/kfp/dsl/_container_op.py | 3 ++- sdk/python/tests/dsl/container_op_tests.py | 11 ++++++++++- 2 files changed, 12 insertions(+), 2 deletions(-) diff --git a/sdk/python/kfp/dsl/_container_op.py b/sdk/python/kfp/dsl/_container_op.py index e5245cb4b6c..30490b3e6f8 100644 --- a/sdk/python/kfp/dsl/_container_op.py +++ b/sdk/python/kfp/dsl/_container_op.py @@ -1182,7 +1182,8 @@ def add_pvolumes(self, self.dependent_names.extend(pvolume.dependent_names) else: pvolume = PipelineVolume(volume=pvolume) - self.pvolumes[mount_path] = pvolume.after(self) + pvolume = pvolume.after(self) + self.pvolumes[mount_path] = pvolume self.add_volume(pvolume) self._container.add_volume_mount(V1VolumeMount( name=pvolume.name, diff --git a/sdk/python/tests/dsl/container_op_tests.py b/sdk/python/tests/dsl/container_op_tests.py index afdd9f4a65c..f3078a16d3a 100644 --- a/sdk/python/tests/dsl/container_op_tests.py +++ b/sdk/python/tests/dsl/container_op_tests.py @@ -17,7 +17,7 @@ from kubernetes.client.models import V1EnvVar, V1VolumeMount import kfp -from kfp.dsl import ContainerOp, UserContainer, Sidecar +from kfp.dsl import ContainerOp, UserContainer, Sidecar, PipelineVolume class TestContainerOp(unittest.TestCase): @@ -89,3 +89,12 @@ def test_deprecation_warnings(self): op.add_volume_mount(V1VolumeMount( mount_path='/secret/gcp-credentials', name='gcp-credentials')) + + + def test_add_pvolumes(self): + pvolume = PipelineVolume(pvc='test') + op = ContainerOp(name='op1', image='image', pvolumes={'/mnt': pvolume}) + + self.assertEqual(pvolume.dependent_names, []) + self.assertEqual(op.pvolume.dependent_names, [op.name]) + self.assertEqual(op.volumes[0].dependent_names, [op.name])