code: def testWorkersDontExceedMaxStreamAssignments(
      self, worker_max_concurrent_snapshots
  ):
      num_workers = 2
      num_snapshots = 10
      cluster = data_service_test_base.TestCluster(
          num_workers=num_workers,
          worker_max_concurrent_snapshots=worker_max_concurrent_snapshots,
      )
      snapshot_dir = data_service_test_base.TempDir()
      paths = []
      for i in range(num_snapshots):
          paths.append(f"{snapshot_dir.full_path}_{i}")
          self.evaluate(
              distributed_save_op.distributed_save(
                  dataset_ops.Dataset.range(5000),
                  paths[i],
                  cluster.dispatcher_address(),
              )
          )
  
      # A mapping of worker idx to max active assignments observed at any time.
      max_assignments = collections.defaultdict(int)
  
      def get_assignments_and_update_max_assignments():
          assignments = get_stream_assignments(
              cluster, num_workers, paths, block=False, active_only=True
          )
          for worker_idx, worker_assignments in assignments.items():
              max_assignments[worker_idx] = max(
                  max_assignments[worker_idx], len(worker_assignments)
              )
          return assignments
  
      # Blocks until each worker has at least the max expected active assignments.
      while True:
          assignments = get_assignments_and_update_max_assignments()
          all_workers_have_assignments = len(assignments) == num_workers
          each_worker_has_enough_assignments = all(
              [
                  len(per_worker_assignments)
                  >= worker_max_concurrent_snapshots
                  for per_worker_assignments in assignments.values()
              ]
          )
          if (
              all_workers_have_assignments
              and each_worker_has_enough_assignments
          ):
              break
          time.sleep(0.1)
  
      cluster.restart_dispatcher()
      wait_for_snapshot(
          paths, get_assignments_and_update_max_assignments
      )
      self.assertValuesEqual(
          list(max_assignments.values()),
          [worker_max_concurrent_snapshots] * num_workers,
      )
  
imports: import collections
  import os
  import pathlib
  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.framework import combinations
  from tensorflow.python.framework import errors
  from tensorflow.python.platform import test
label_desc: Use image node_id 12 for calling the SnapshotFtTest obj's underlying member method code with example usage: obj.testWorkersDontExceedMaxStreamAssignments(worker_max_concurrent_snapshots) and returns: assignments