code: def _traverse(
      node_map: Mapping[str, node_proto_view.NodeProtoView],
      direction: _Direction,
      start_nodes: Collection[str],
  ) -> Set[str]:
      """Traverses a DAG from start_nodes, either upstream or downstream.
  
      Args:
        node_map: Mapping of node_id to nodes.
        direction: _Direction.UPSTREAM or _Direction.DOWNSTREAM.
        start_nodes: node_ids to start from.
  
      Returns:
        Set of node_ids visited by this traversal.
      """
      result = set()
      stack = []
      for start_node in start_nodes:
          # Depth-first traversal
          stack.append(start_node)
          while stack:
              current_node_id = stack.pop()
              if current_node_id in result:
                  continue
              result.add(current_node_id)
              if direction == _Direction.UPSTREAM:
                  stack.extend(node_map[current_node_id].upstream_nodes)
              elif direction == _Direction.DOWNSTREAM:
                  stack.extend(
                      node_map[current_node_id].downstream_nodes
                  )
      return result
  
imports: import collections
  import concurrent.futures
  import enum
  from typing import Collection, Dict, Final, List, Mapping, Optional, OrderedDict, Sequence, Set, Tuple
  import uuid
  from absl import logging
  from tfx.dsl.compiler import compiler_utils
  from tfx.dsl.compiler import constants
  from tfx.orchestration import metadata
  from tfx.orchestration import node_proto_view
  from tfx.orchestration.portable import execution_publish_utils
  from tfx.orchestration.portable.mlmd import context_lib
  from tfx.orchestration.portable.mlmd import execution_lib
  from tfx.proto.orchestration import pipeline_pb2
  from ml_metadata.proto import metadata_store_pb2
label_desc: Use image node_id 14 for calling a global function with example usage: _traverse(node_map, direction, start_nodes) and returns: result