base_node.py

# Copyright 2019 Google LLC. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""Base class for TFX nodes."""

import abc
import copy
from typing import Any, Dict, Optional, Type

from tfx.dsl.components.base import base_driver
from tfx.dsl.components.base import base_executor
from tfx.dsl.components.base import executor_spec as executor_spec_module
from tfx.dsl.context_managers import dsl_context_registry
from tfx.dsl.experimental.node_execution_options import utils
from tfx.utils import deprecation_utils
from tfx.utils import doc_controls
from tfx.utils import json_utils
from tfx.utils import name_utils
import typing_extensions


def _abstract_property() -> Any:
  """Returns an abstract property for use in an ABC abstract class."""
  return abc.abstractmethod(lambda: None)


class BaseNode(json_utils.Jsonable, abc.ABC):
  """Base class for a node in TFX pipeline."""

  def __new__(cls, *args, **kwargs):
    # Record invocation details for tracing. No backwards-compatibility
    # guarantees; for TFX-internal use only.
    result = super(BaseNode, cls).__new__(cls)
    result._CONSTRUCT_CLS = cls
    result._CONSTRUCT_ARGS = args
    result._CONSTRUCT_KWARGS = kwargs
    return result

  def __init__(
      self,
      executor_spec: Optional[executor_spec_module.ExecutorSpec] = None,
      driver_class: Optional[Type[base_driver.BaseDriver]] = None,
  ):
    """Initialize a node.

    Args:
      executor_spec: Optional instance of executor_spec.ExecutorSpec which
        describes how to execute this node (optional, defaults to an empty
        executor indicates no-op.
      driver_class: Optional subclass of base_driver.BaseDriver as a custom
        driver for this node (optional, defaults to base_driver.BaseDriver).
        Nodes usually use the default driver class, but may override it.
    """
    if executor_spec is None:
      executor_spec = executor_spec_module.ExecutorClassSpec(
          base_executor.EmptyExecutor)
    if driver_class is None:
      driver_class = base_driver.BaseDriver
    self.executor_spec = executor_spec
    self.driver_class = driver_class
    self._upstream_nodes = set()
    self._downstream_nodes = set()
    self._id = None
    self._node_execution_options: Optional[utils.NodeExecutionOptions] = None
    dsl_context_registry.get().put_node(self)

  @doc_controls.do_not_doc_in_subclasses
  def to_json_dict(self) -> Dict[str, Any]:
    """Convert from an object to a JSON serializable dictionary."""
    return dict((k, v)
                for k, v in self.__dict__.items()
                if k not in ['_upstream_nodes', '_downstream_nodes'])

  @classmethod
  @doc_controls.do_not_doc_in_subclasses
  def get_class_type(cls) -> str:
    nondeprecated_class = deprecation_utils.get_first_nondeprecated_class(cls)
    # TODO(b/221166027): Turn strict_check=True once failing tests are fixed.
    return name_utils.get_full_name(nondeprecated_class, strict_check=False)

  @property
  @doc_controls.do_not_doc_in_subclasses
  def type(self) -> str:
    return self.__class__.get_class_type()

  @property
  @deprecation_utils.deprecated(None,
                                'component_type is deprecated, use type instead'
                               )
  @doc_controls.do_not_doc_in_subclasses
  def component_type(self) -> str:
    return self.type

  @property
  @doc_controls.do_not_doc_in_subclasses
  def id(self) -> str:
    """Node id, unique across all TFX nodes in a pipeline.

    If `id` is set by the user, return it directly.
    Otherwise, return <node_class_name>.

    Returns:
      node id.
    """
    if self._id:
      return self._id
    node_class = deprecation_utils.get_first_nondeprecated_class(self.__class__)
    return node_class.__name__

  @property
  @deprecation_utils.deprecated(None,
                                'component_id is deprecated, use id instead')
  @doc_controls.do_not_doc_in_subclasses
  def component_id(self) -> str:
    return self.id

  @id.setter
  @doc_controls.do_not_doc_in_subclasses
  def id(self, id: str) -> None:  # noqa: A002
    self._id = id

  # TODO(kmonte): Update this to Self once we're on 3.11 everywhere
  @doc_controls.do_not_doc_in_subclasses
  def with_id(self, id: str) -> typing_extensions.Self:  # noqa: A002
    self._id = id
    return self

  @property
  @abc.abstractmethod
  def inputs(self) -> Dict[str, Any]:
    pass

  @property
  @abc.abstractmethod
  def outputs(self) -> Dict[str, Any]:
    pass

  @property
  @abc.abstractmethod
  def exec_properties(self) -> Dict[str, Any]:
    pass

  @property
  @doc_controls.do_not_doc_in_subclasses
  def upstream_nodes(self):
    return self._upstream_nodes

  @property
  @doc_controls.do_not_doc_in_subclasses
  def node_execution_options(self) -> Optional[utils.NodeExecutionOptions]:
    return self._node_execution_options

  @node_execution_options.setter
  @doc_controls.do_not_doc_in_subclasses
  def node_execution_options(
      self,
      node_execution_options: utils.NodeExecutionOptions
  ):
    self._node_execution_options = copy.deepcopy(node_execution_options)

  # TODO(kmonte): Update this to Self once we're on 3.11 everywhere
  def with_node_execution_options(
      self, node_execution_options: utils.NodeExecutionOptions
  ) -> typing_extensions.Self:
    self.node_execution_options = node_execution_options
    return self

  @doc_controls.do_not_doc_in_subclasses
  def add_upstream_node(self, upstream_node):
    """Experimental: Add another component that must run before this one.

    This method enables task-based dependencies by enforcing execution order for
    synchronous pipelines on supported platforms. Currently, the supported
    platforms are Airflow, Beam, and Kubeflow Pipelines.

    Note that this API call should be considered experimental, and may not work
    with asynchronous pipelines, sub-pipelines and pipelines with conditional
    nodes. We also recommend relying on data for capturing dependencies where
    possible to ensure data lineage is fully captured within MLMD.

    It is symmetric with `add_downstream_node`.

    Args:
      upstream_node: a component that must run before this node.
    """
    self._upstream_nodes.add(upstream_node)
    if self not in upstream_node.downstream_nodes:
      upstream_node.add_downstream_node(self)

  @doc_controls.do_not_doc_in_subclasses
  def add_upstream_nodes(self, upstream_nodes):
    """Experimental: Add components that must run before this one.

    This method enables task-based dependencies by enforcing execution order for
    synchronous pipelines on supported platforms. Currently, the supported
    platforms are Airflow, Beam, and Kubeflow Pipelines.

    Note that this API call should be considered experimental, and may not work
    with asynchronous pipelines, sub-pipelines and pipelines with conditional
    nodes. We also recommend relying on data for capturing dependencies where
    possible to ensure data lineage is fully captured within MLMD.


    Args:
      upstream_nodes: a list of components that must run before this node.
    """
    self._upstream_nodes.update(upstream_nodes)
    for upstream_node in upstream_nodes:
      if self not in upstream_node.downstream_nodes:
        upstream_node.add_downstream_node(self)

  @doc_controls.do_not_doc_in_subclasses
  def remove_upstream_node(self, upstream_node):
    self._upstream_nodes.remove(upstream_node)
    if self in upstream_node.downstream_nodes:
      upstream_node.remove_downstream_node(self)

  @property
  @doc_controls.do_not_doc_in_subclasses
  def downstream_nodes(self):
    return self._downstream_nodes

  @doc_controls.do_not_doc_in_subclasses
  def add_downstream_node(self, downstream_node):
    """Experimental: Add another component that must run after this one.

    This method enables task-based dependencies by enforcing execution order for
    synchronous pipelines on supported platforms. Currently, the supported
    platforms are Airflow, Beam, and Kubeflow Pipelines.

    Note that this API call should be considered experimental, and may not work
    with asynchronous pipelines, sub-pipelines and pipelines with conditional
    nodes. We also recommend relying on data for capturing dependencies where
    possible to ensure data lineage is fully captured within MLMD.

    It is symmetric with `add_upstream_node`.

    Args:
      downstream_node: a component that must run after this node.
    """
    self._downstream_nodes.add(downstream_node)
    if self not in downstream_node.upstream_nodes:
      downstream_node.add_upstream_node(self)

  @doc_controls.do_not_doc_in_subclasses
  def add_downstream_nodes(self, downstream_nodes):
    """Experimental: Add another component that must run after this one.

    This method enables task-based dependencies by enforcing execution order for
    synchronous pipelines on supported platforms. Currently, the supported
    platforms are Airflow, Beam, and Kubeflow Pipelines.

    Note that this API call should be considered experimental, and may not work
    with asynchronous pipelines, sub-pipelines and pipelines with conditional
    nodes. We also recommend relying on data for capturing dependencies where
    possible to ensure data lineage is fully captured within MLMD.

    It is symmetric with `add_upstream_nodes`.

    Args:
      downstream_nodes: a list of components that must run after this node.
    """
    self._downstream_nodes.update(downstream_nodes)
    for downstream_node in downstream_nodes:
      if self not in downstream_node.upstream_nodes:
        downstream_node.add_upstream_node(self)

  @doc_controls.do_not_doc_in_subclasses
  def remove_downstream_node(self, downstream_node):
    self._downstream_nodes.remove(downstream_node)
    if self in downstream_node.upstream_nodes:
      downstream_node.remove_upstream_node(self)