<commit_before># Copyright 2022 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.
"""Module for compiler context classes definition."""

from __future__ import annotations

import collections
from typing import Optional, List, Set, Iterable, Dict

from tfx.dsl.compiler import compiler_utils
from tfx.dsl.components.base import base_node
from tfx.dsl.context_managers import dsl_context_registry
from tfx.dsl.control_flow import for_each_internal
from tfx.dsl.experimental.conditionals import conditional
from tfx.dsl.input_resolution import resolver_op
from tfx.orchestration import pipeline
from tfx.proto.orchestration import pipeline_pb2
from tfx.types import channel as channel_types


class PipelineContext:
  """Encapsulates resources needed to compile a pipeline.

  `PipelineContext` is for compiling a single pipeline. If the pipeline is
  compiling a sub-pipeline (a `Pipeline` node as a `pipeline.components`) then
  another `PipelineContext` would be created and used.

  Use it as a first positional argument to the compiler extension functions as
  if you are extending a class and this is a self object.
  """

  def __init__(self,
               tfx_pipeline: pipeline.Pipeline,
               parent: Optional[PipelineContext] = None):
    self.pipeline = tfx_pipeline

    self.pipeline_info = self.pipeline.pipeline_info
    self.execution_mode = compiler_utils.resolve_execution_mode(self.pipeline)
    self.dsl_context_registry = self.pipeline.dsl_context_registry
    self.parent = parent or NullPipelineContext(child=self)

    # Stores channels available in the current pipeline scope.
    # Mapping from Channel object to compiled Channel proto.
    self.channels = dict()

    # Node ID -> NodeContext
    self._node_contexts: Dict[str, NodeContext] = {}

    # CondContext -> Context ID
    self._conditional_ids: Dict[conditional.CondContext, str] = {}

    self._pipeline_nodes_by_id = {}
    self._topological_order = {}
    self._implicit_upstream_nodes = collections.defaultdict(set)
    self._implicit_downstream_nodes = collections.defaultdict(set)
    for i, node in enumerate(tfx_pipeline.components, start=1):
      self._pipeline_nodes_by_id[node.id] = node
      self._topological_order[node.id] = i
      self._collect_conditional_dependency(node)

  @property
  def is_root(self) -> bool:
    return isinstance(self.parent, NullPipelineContext)

  @property
  def is_subpipeline(self) -> bool:
    return self.parent and self.parent.pipeline is not None

  @property
  def parent_pipelines(self) -> List[pipeline.Pipeline]:
    """Parent pipelines, in the order of outer -> inner."""
    result = []
    ctx = self.parent
    while ctx and ctx.pipeline:
      result.append(ctx.pipeline)
      ctx = ctx.parent
    return result[::-1]

  def _add_implicit_dependency(self, parent_id: str, child_id: str) -> None:
    self._implicit_upstream_nodes[child_id].add(parent_id)
    self._implicit_downstream_nodes[parent_id].add(child_id)

  def _collect_conditional_dependency(self, here: base_node.BaseNode) -> None:
    for predicate in conditional.get_predicates(here,
                                                self.dsl_context_registry):
      for chnl in predicate.dependent_channels():
        if isinstance(chnl, channel_types.OutputChannel):
          self._add_implicit_dependency(chnl.producer_component_id, here.id)

  def topologically_sorted(
      self, tfx_nodes: Iterable[base_node.BaseNode],
  ) -> List[base_node.BaseNode]:
    return sorted(tfx_nodes, key=lambda node: self._topological_order[node.id])

  @property
  def is_sync_mode(self):
    return self.execution_mode == pipeline_pb2.Pipeline.SYNC

  @property
  def is_async_mode(self):
    return self.execution_mode == pipeline_pb2.Pipeline.ASYNC

  @property
  def pipeline_node_ids(self) -> Set[str]:
    return set(self._pipeline_nodes_by_id.keys())

  def implicit_upstream_nodes(
      self, here: base_node.BaseNode) -> List[base_node.BaseNode]:
    return [
        self._pipeline_nodes_by_id[node_id]
        for node_id in self._implicit_upstream_nodes[here.id]
    ]

  def implicit_downstream_nodes(
      self, here: base_node.BaseNode) -> List[base_node.BaseNode]:
    return [
        self._pipeline_nodes_by_id[node_id]
        for node_id in self._implicit_downstream_nodes[here.id]
    ]

  def get_node_context(self, tfx_node: base_node.BaseNode) -> NodeContext:
    if tfx_node.id not in self._node_contexts:
      self._node_contexts[tfx_node.id] = NodeContext(tfx_node)
    return self._node_contexts[tfx_node.id]

  def get_conditional_id(self, cond_context: conditional.CondContext) -> str:
    if cond_context not in self._conditional_ids:
      cond_id = f'cond_{len(self._conditional_ids) + 1}'
      self._conditional_ids[cond_context] = cond_id
    return self._conditional_ids[cond_context]


class NullPipelineContext(PipelineContext):
  """PipelineContext without pipeline.

  NullPipelineContext is the parent PipelineContext for the root level pipeline.

  NullPipelineContext is used when compiling a pipeline begin node of the
  root level pipeline, since pipeline begin node is compiled as if compiling a
  pipeline as a node from the parent pipeline, which does not exist for the
  root level pipeline.
  """

  def __init__(self, *, child: PipelineContext):
    # pylint: disable=super-init-not-called
    # Doesn't invoke super().__init__() as it is not a typical PipelineContext
    # with a pipeline.
    self.pipeline = None
    self.pipeline_info = None
    self.dsl_context_registry = dsl_context_registry.DslContextRegistry()
    self.dsl_context_registry.put_node(child.pipeline)
    self.parent = None
    self.channels = {}

    # Node ID -> NodeContext
    self._node_contexts: Dict[str, NodeContext] = {}


class NodeContext:
  """Per-node context for pipeline compilation."""

  def __init__(self, tfx_node: base_node.BaseNode):
    self._input_key_by_channel = {
        channel: input_key
        for input_key, channel in tfx_node.inputs.items()<commit_msg>add _implicit_downstream_nodes to NullPipelineContext<commit_after># Copyright 2022 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.
"""Module for compiler context classes definition."""

from __future__ import annotations

import collections
from typing import Optional, List, Set, Iterable, Dict

from tfx.dsl.compiler import compiler_utils
from tfx.dsl.components.base import base_node
from tfx.dsl.context_managers import dsl_context_registry
from tfx.dsl.control_flow import for_each_internal
from tfx.dsl.experimental.conditionals import conditional
from tfx.dsl.input_resolution import resolver_op
from tfx.orchestration import pipeline
from tfx.proto.orchestration import pipeline_pb2
from tfx.types import channel as channel_types


class PipelineContext:
  """Encapsulates resources needed to compile a pipeline.

  `PipelineContext` is for compiling a single pipeline. If the pipeline is
  compiling a sub-pipeline (a `Pipeline` node as a `pipeline.components`) then
  another `PipelineContext` would be created and used.

  Use it as a first positional argument to the compiler extension functions as
  if you are extending a class and this is a self object.
  """

  def __init__(self,
               tfx_pipeline: pipeline.Pipeline,
               parent: Optional[PipelineContext] = None):
    self.pipeline = tfx_pipeline

    self.pipeline_info = self.pipeline.pipeline_info
    self.execution_mode = compiler_utils.resolve_execution_mode(self.pipeline)
    self.dsl_context_registry = self.pipeline.dsl_context_registry
    self.parent = parent or NullPipelineContext(child=self)

    # Stores channels available in the current pipeline scope.
    # Mapping from Channel object to compiled Channel proto.
    self.channels = dict()

    # Node ID -> NodeContext
    self._node_contexts: Dict[str, NodeContext] = {}

    # CondContext -> Context ID
    self._conditional_ids: Dict[conditional.CondContext, str] = {}

    self._pipeline_nodes_by_id = {}
    self._topological_order = {}
    self._implicit_upstream_nodes = collections.defaultdict(set)
    self._implicit_downstream_nodes = collections.defaultdict(set)
    for i, node in enumerate(tfx_pipeline.components, start=1):
      self._pipeline_nodes_by_id[node.id] = node
      self._topological_order[node.id] = i
      self._collect_conditional_dependency(node)

  @property
  def is_root(self) -> bool:
    return isinstance(self.parent, NullPipelineContext)

  @property
  def is_subpipeline(self) -> bool:
    return self.parent and self.parent.pipeline is not None

  @property
  def parent_pipelines(self) -> List[pipeline.Pipeline]:
    """Parent pipelines, in the order of outer -> inner."""
    result = []
    ctx = self.parent
    while ctx and ctx.pipeline:
      result.append(ctx.pipeline)
      ctx = ctx.parent
    return result[::-1]

  def _add_implicit_dependency(self, parent_id: str, child_id: str) -> None:
    self._implicit_upstream_nodes[child_id].add(parent_id)
    self._implicit_downstream_nodes[parent_id].add(child_id)

  def _collect_conditional_dependency(self, here: base_node.BaseNode) -> None:
    for predicate in conditional.get_predicates(here,
                                                self.dsl_context_registry):
      for chnl in predicate.dependent_channels():
        if isinstance(chnl, channel_types.OutputChannel):
          self._add_implicit_dependency(chnl.producer_component_id, here.id)

  def topologically_sorted(
      self, tfx_nodes: Iterable[base_node.BaseNode],
  ) -> List[base_node.BaseNode]:
    return sorted(tfx_nodes, key=lambda node: self._topological_order[node.id])

  @property
  def is_sync_mode(self):
    return self.execution_mode == pipeline_pb2.Pipeline.SYNC

  @property
  def is_async_mode(self):
    return self.execution_mode == pipeline_pb2.Pipeline.ASYNC

  @property
  def pipeline_node_ids(self) -> Set[str]:
    return set(self._pipeline_nodes_by_id.keys())

  def implicit_upstream_nodes(
      self, here: base_node.BaseNode) -> List[base_node.BaseNode]:
    return [
        self._pipeline_nodes_by_id[node_id]
        for node_id in self._implicit_upstream_nodes[here.id]
    ]

  def implicit_downstream_nodes(
      self, here: base_node.BaseNode) -> List[base_node.BaseNode]:
    return [
        self._pipeline_nodes_by_id[node_id]
        for node_id in self._implicit_downstream_nodes[here.id]
    ]

  def get_node_context(self, tfx_node: base_node.BaseNode) -> NodeContext:
    if tfx_node.id not in self._node_contexts:
      self._node_contexts[tfx_node.id] = NodeContext(tfx_node)
    return self._node_contexts[tfx_node.id]

  def get_conditional_id(self, cond_context: conditional.CondContext) -> str:
    if cond_context not in self._conditional_ids:
      cond_id = f'cond_{len(self._conditional_ids) + 1}'
      self._conditional_ids[cond_context] = cond_id
    return self._conditional_ids[cond_context]


class NullPipelineContext(PipelineContext):
  """PipelineContext without pipeline.

  NullPipelineContext is the parent PipelineContext for the root level pipeline.

  NullPipelineContext is used when compiling a pipeline begin node of the
  root level pipeline, since pipeline begin node is compiled as if compiling a
  pipeline as a node from the parent pipeline, which does not exist for the
  root level pipeline.
  """

  def __init__(self, *, child: PipelineContext):
    # pylint: disable=super-init-not-called
    # Doesn't invoke super().__init__() as it is not a typical PipelineContext
    # with a pipeline.
    self.pipeline = None
    self.pipeline_info = None
    self.dsl_context_registry = dsl_context_registry.DslContextRegistry()
    self.dsl_context_registry.put_node(child.pipeline)
    self.parent = None
    self.channels = {}
    self._implicit_downstream_nodes = collections.defaultdict(set)

    # Node ID -> NodeContext
    self._node_contexts: Dict[str, NodeContext] = {}


class NodeContext:
  """Per-node context for pipeline compilation."""

  def __init__(self, tfx_node: base_node.BaseNode):
    self._input_key_by_channel = {
        channel: input_key
        for input_key, channel in tfx_node.inputs.items()
{'fraction_non_alphanumeric': 0.08127315339835185, 'fraction_numerical': 0.002116882135026839, 'mean_word_length': 3.9935824839562097, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 33, 'https://': 0, 'lorem ipsum': 0, 'www.': 2, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 4, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '6067320', 'n_tokens_mistral': 3950, 'n_tokens_neox': 3865, 'n_words': 1076}