code: def enforce_rescale(
      cluster,
      scale,
      max_attempts=DEFAULT_MAX_ATTEMPT,
      wait_time=DEFAULT_WAIT_TIME,
  ):
      cluster.scale(scale)
      attempt = 0
      ready = len(cluster.workers) == scale
      while (attempt < max_attempts) and not ready:
          time.sleep(wait_time)
          ready = len(cluster.workers) == scale
          attempt += 1
      assert ready, "Unable to rescale cluster to {}".format(scale)
  
file_path: cugraph/python/cugraph/cugraph/tests/utils/mg_context.py
imports: import time
  import pytest
  import cugraph.dask.comms
  from dask.distributed import Client
  from dask_cuda import LocalCUDACluster
  from cugraph.dask.common.mg_utils import get_visible_devices
label_desc: Use image node_id 2 for calling a global function with example usage: enforce_rescale(cluster, scale, max_attempts, wait_time) without return types