[sdk] Wrong DSL component and pipeline return types

[AndersBennedsgaard]

Environment

• KFP version: 2.1.0 (not relevant)
• KFP SDK version: 2.7.0
• All dependencies version:

  kfp                      2.7.0
  kfp-pipeline-spec        0.3.0
  kfp-server-api           2.0.5
  

Steps to reproduce

When creating a component and pipeline with a return value using the DSL, you get a type error since the decorator changes the return type.
For example, a simple add(a: float, b: float) -> float component will always cause a type error, since the decorator makes the function return a dsl.PipelineTask or Unknown instead. Pipelines also always return a Callable or Unknown:

from kfp import dsl, local

local.init(runner=local.DockerRunner())

@dsl.component(base_image="python:3.8-slim")
def add(a: float, b: float) -> float:
    return a + b

@dsl.pipeline()  # call the decorator with (), otherwise the pipeline type is Unknown
def add_pipeline(a: float, b: float) -> float:
    t1 = add(a=a, b=b)
    return t1.output
    #         ^^^ Cannot access member "output" for type "float"

p = add_pipeline(a=1.0, b=2.0)
assert p.output == 3.0
    #    ^^^ Cannot access member "output" for type "float"

Expected result

Components and pipelines should return the output that the user specifies, and not have it wrapped in an output class.

Alternatively, in order to still return the custom output classes and not introduce breaking changes and a large refactor in the code base, it is possible to make use of TypeVar and Generic magic. Based on help from https://stackoverflow.com/questions/65936408/typing-function-when-decorator-change-generic-return-type, I've created this simple proof of concept:

from typing import Callable, Generic, ParamSpec, TypeVar

T = TypeVar("T")

class PipelineTask(Generic[T]):
    def __init__(self, item: T):
        self.item: T = item

    @property
    def output(self) -> T:
        return self.item

P = ParamSpec("P")
R = TypeVar("R")

# does not support arguments, but it can be extended for this
def my_component(f: Callable[P, R]) -> Callable[P, PipelineTask[R]]:
    def wrapper(*args: P.args, **kwargs: P.kwargs) -> PipelineTask[R]:
        return PipelineTask(f(*args, **kwargs))

    return wrapper

@my_component
def add(a: float, b: float) -> float:
    return a + b

out = add(1.0, 2.0)  # type PipelineTask[float]

assert out.output == 3.0  # no type errors

Materials and Reference

• https://stackoverflow.com/questions/65936408/typing-function-when-decorator-change-generic-return-type

---

Impacted by this bug? Give it a 👍.