code: def _load_distributed_snapshot_v2(
      path: str, reader_func=None
  ) -> dataset_ops.Dataset:
      """Load a distributed snapshot using the updated loading algorithm.
  
      The new version allows the load job to read the snapshot while it is being
      written.
  
      TODO(b/297930782): Merge this into `_load` when it's ready. Currently, this is
      for testing only.
  
      Args:
        path: Base path of the snapshot.
        reader_func: Optional. A function to control how to read data from shards.
          If present, the function will be traced and executed as graph computation.
  
      Returns:
        The loaded dataset.
      """
  
      if not reader_func:
          reader_func = lambda datasets: datasets.interleave(  # pylint:disable=g-long-lambda
              lambda x: x,
              cycle_length=multiprocessing.cpu_count(),
              num_parallel_calls=dataset_ops.AUTOTUNE,
          )
  
      metadata = _load_distributed_snapshot_metadata(path)
      while not metadata:
          time.sleep(2)
          metadata = _load_distributed_snapshot_metadata(path)
  
      dataset = _ListSnapshotChunksDataset(path)
      dataset = dataset.map(
          lambda chunk_file: _SnapshotChunkDataset(  # pylint:disable=g-long-lambda
              chunk_file,
              element_spec=_parse_element_spec(metadata.element_spec),
              compression=metadata.compression,
          )
      )
      return reader_func(dataset)
  
imports: import multiprocessing
  import os
  import time
  from typing import Optional
  from google.protobuf import message
  from google.protobuf import text_format
  from tensorflow.core.protobuf import snapshot_pb2
  from tensorflow.python.data.experimental.service import _pywrap_snapshot_utils
  from tensorflow.python.data.ops import dataset_ops
  from tensorflow.python.data.ops import structured_function
  from tensorflow.python.framework import dtypes
  from tensorflow.python.framework import errors
  from tensorflow.python.framework import tensor_spec
  from tensorflow.python.ops import gen_experimental_dataset_ops
  from tensorflow.python.platform import gfile
  from tensorflow.python.saved_model import nested_structure_coder
label_desc: Use image node_id 4 for calling a global function with example usage: _load_distributed_snapshot_v2(path, reader_func) and returns: reader_func