Skip to content

Commit

Permalink
test: Add ParallelFor compile test over single artifact. (kubeflow#10531
Browse files Browse the repository at this point in the history
)

* .

* formatting
  • Loading branch information
KevinGrantLee authored and petethegreat committed Mar 29, 2024
1 parent 4a583c7 commit 9e207c7
Show file tree
Hide file tree
Showing 3 changed files with 16 additions and 2 deletions.
14 changes: 14 additions & 0 deletions sdk/python/kfp/compiler/compiler_test.py
Expand Up @@ -838,6 +838,20 @@ def my_pipeline():
with dsl.ParallelFor(items=single_param_task.output) as item:
print_and_return(text=item)

def test_cannot_compile_parallel_for_with_single_artifact(self):

with self.assertRaisesRegex(
ValueError,
r'Cannot iterate over a single artifact using `dsl\.ParallelFor`\. Expected a list of artifacts as argument to `items`\.'
):

@dsl.pipeline
def my_pipeline():
single_artifact_task = print_and_return_as_artifact(
text='string')
with dsl.ParallelFor(items=single_artifact_task.output) as item:
print_artifact(a=item)

def test_pipeline_in_pipeline(self):

@dsl.pipeline(name='graph-component')
Expand Down
2 changes: 1 addition & 1 deletion sdk/python/kfp/dsl/for_loop.py
Expand Up @@ -286,7 +286,7 @@ def from_pipeline_channel(
object."""
if not channel.is_artifact_list:
raise ValueError(
'Cannot iterate over a single Artifact using `dsl.ParallelFor`. Expected a list of Artifacts as argument to `items`.'
'Cannot iterate over a single artifact using `dsl.ParallelFor`. Expected a list of artifacts as argument to `items`.'
)
return LoopArtifactArgument(
items=channel,
Expand Down
2 changes: 1 addition & 1 deletion sdk/python/kfp/dsl/for_loop_test.py
Expand Up @@ -207,7 +207,7 @@ def test_loop_artifact_argument_from_single_pipeline_channel_raises_error(
self, channel):
with self.assertRaisesRegex(
ValueError,
r'Cannot iterate over a single Artifact using `dsl\.ParallelFor`\. Expected a list of Artifacts as argument to `items`\.'
r'Cannot iterate over a single artifact using `dsl\.ParallelFor`\. Expected a list of artifacts as argument to `items`\.'
):
loop_argument = for_loop.LoopArtifactArgument.from_pipeline_channel(
channel)
Expand Down

0 comments on commit 9e207c7

Please sign in to comment.