code: def get_stream_assignment(
      cluster, worker_idx, path, block=True, active_only=False
  ):
      while True:
          for progress in cluster.workers[
              worker_idx
          ].snapshot_task_progresses():
              if (
                  progress.snapshot_task_base_path.decode() == path
                  and not (active_only and progress.completed)
              ):
                  return progress.snapshot_task_stream_index
          if not block:
              break
          time.sleep(0.1)
  
imports: import collections
  import os
  import pathlib
  import shutil
  import tempfile
  import time
  from absl.testing import parameterized
  from tensorflow.python.data.experimental.kernel_tests.service import test_base
  from tensorflow.python.data.experimental.ops import distributed_save_op
  from tensorflow.python.data.experimental.service import _pywrap_snapshot_utils
  from tensorflow.python.data.kernel_tests import test_base
  from tensorflow.python.data.ops import dataset_ops
  from tensorflow.python.data.ops import test_mode
  from tensorflow.python.framework import combinations
  from tensorflow.python.framework import errors
  from tensorflow.python.platform import test
label_desc: Use image node_id 2 for calling a global function with example usage: get_stream_assignment(cluster, worker_idx, path, block, active_only) and returns: progress