Join the SIG TFX-Addons community and help make TFX even better!


Decorator: creates a component from a typehint-annotated Python function.

This decorator creates a component based on typehint annotations specified for the arguments and return value for a Python function. Specifically, function arguments can be annotated with the following types and associated semantics:

  • Parameter[T] where T is int, float, str, or bytes: indicates that a primitive type execution parameter, whose value is known at pipeline construction time, will be passed for this argument. These parameters will be recorded in ML Metadata as part of the component's execution record. Can be an optional argument.
  • int, float, str, bytes: indicates that a primitive type value will be passed for this argument. This value is tracked as an Integer, Float String or Bytes artifact (see tfx.types.standard_artifacts) whose value is read and passed into the given Python component function. Can be an optional argument.
  • InputArtifact[ArtifactType]: indicates that an input artifact object of type ArtifactType (deriving from tfx.types.Artifact) will be passed for this argument. This artifact is intended to be consumed as an input by this component (possibly reading from the path specified by its .uri). Can be an optional argument by specifying a default value of None.
  • OutputArtifact[ArtifactType]: indicates that an output artifact object of type ArtifactType (deriving from tfx.types.Artifact) will be passed for this argument. This artifact is intended to be emitted as an output by this component (and written to the path specified by its .uri). Cannot be an optional argument.

The return value typehint should be either empty or None, in the case of a component function that has no return values, or an instance of OutputDict(key_1=type_1, ...), where each key maps to a given type (each type is a primitive value type, i.e. int, float, str or bytes), to indicate that the return value is a dictionary with specified keys and value types.

Note that output artifacts should not be included in the return value typehint; they should be included as OutputArtifact annotations in the function inputs, as described above.

The function to which this decorator is applied must be at the top level of its Python module (it may not be defined within nested classes or function closures).

This is example usage of component definition using this decorator:

from tfx.dsl.components.base.annotations import OutputDict
from tfx.dsl.components.base.annotations import
from tfx.dsl.components.base.annotations import
from tfx.dsl.components.base.annotations import
from tfx.dsl.components.base.decorators import component
from tfx.types.standard_artifacts import Examples
from tfx.types.standard_artifacts import Model

def MyTrainerComponent(
    training_data: InputArtifact[Examples],
    model: OutputArtifact[Model],
    dropout_hyperparameter: float,
    num_iterations: Parameter[int] = 10
    ) -> OutputDict(loss=float, accuracy=float):
  '''My simple trainer component.'''

  records = read_examples(training_data.uri)
  model_obj = train_model(records, num_iterations, dropout_hyperparameter)

  return {
    'loss': model_obj.loss,
    'accuracy': model_obj.accuracy

Example:usage in a pipeline graph definition:
# ...
trainer = MyTrainerComponent(
pusher = Pusher(model=trainer.outputs['model'])
# ...

Experimental: no backwards compatibility guarantees.

func Typehint-annotated component executor function.

base_component.BaseComponent subclass for the given component executor function.

EnvironmentError if the current Python interpreter is not Python 3.