Use dsl.Collected inside a dsl.Condition context

Question

I'm trying to use the output of dsl.ParallelFor downstream within a dsl.Condition and I get this error: TypeError: None has type NoneType, but expected one of: int
I know regular output of components can be use in dsl.Condition contexts without issues, and dsl.Collected outside the context works fine, so is there any way to get the best of both worlds and get this to work ?

from typing import List

from kfp import compiler, dsl
from kfp.dsl import Input, Model, Output


@dsl.component
def some_string(k: int, model: Output[Model]):
    model = "x"


@dsl.component
def always_one() -> int:
    return 1


@dsl.component
def component_with_collection(models: Input[List[Model]]) -> str:
    return ".".join([model for model in models])


@dsl.pipeline
def main_pipeline():
    always_one_task = always_one()

    with dsl.ParallelFor(items=[1, 5, 10, 25], parallelism=2) as i:
        some_string_task = some_string(k=i)

    with dsl.Condition(always_one_task.output == 1):
        component_task = component_with_collection(
            models=dsl.Collected(some_string_task.outputs["model"])
        )


compiler.Compiler().compile(main_pipeline, "main_compiled.yaml")

Use dsl.Collected inside a dsl.Condition context

0 Answers0