Python test_util.create_local_cluster函数代码示例

  def testDistributedOOM(self):
    if not test.is_gpu_available():
      return
    ops.reset_default_graph()

    workers, _ = test_util.create_local_cluster(2, 0)

    with ops.device('/job:worker/replica:0/task:0/gpu:0'):
      a = random_ops.random_normal([1, 10000, 20000], name='test_random1')
    with ops.device('/job:worker/replica:0/task:1/gpu:0'):
      b = random_ops.random_normal([30000, 10000, 1], name='test_random2')
      c = a * b

    try:
      with session.Session(workers[1].target) as sess:
        sess.run(c, options=config_pb2.RunOptions(
            report_tensor_allocations_upon_oom=True))
    except Exception as e:  # pylint: disable=broad-except
      exception_str = '%s' % e
      # test_random2 is reported because it's allocated in worker 1.
      self.assertTrue('Current usage from device: '
                      '/job:worker/replica:0/task:1/device:GPU:0, '
                      'allocator: GPU_0_bfc' in exception_str)
      mat = re.search('(.*)GiB from test_random2/RandomStandardNormal',
                      exception_str)
      self.assertGreater(float(mat.group(1)), 0.0)
      # test_random1 is not reported because it's allocated in worker 0.
      mat = re.search('(.*)MiB from test_random1/RandomStandardNormal',
                      exception_str)
      self.assertTrue(mat is None)

  def testImplicitDisposeParallelMapDataset(self):
    # Tests whether a parallel map dataset will be cleaned up correctly when
    # the pipeline does not run it until exhaustion.
    # The pipeline is TensorSliceDataset -> MapDataset(square_3) ->
    # RepeatDataset(None) -> PrefetchDataset(100).
    worker, _ = test_util.create_local_cluster(1, 1)

    components = (np.arange(1000),
                  np.array([[1, 2, 3]]) * np.arange(1000)[:, np.newaxis],
                  np.array(37.0) * np.arange(1000))

    def _map_fn(x, y, z):
      return math_ops.square(x), math_ops.square(y), math_ops.square(z)

    dataset = (
        dataset_ops.Dataset.from_tensor_slices(components).map(_map_fn)
        .repeat(None).prefetch(10000))

    iterator = dataset.make_initializable_iterator()
    init_op = iterator.initializer
    get_next = iterator.get_next()

    with session.Session(worker[0].target) as sess:
      self.evaluate(init_op)
      for _ in range(3):
        sess.run(get_next)

  def testRemoteDeviceInPartitionedCallOp(self):
    workers, _ = test_util.create_local_cluster(2, 0)

    worker0_device = "/job:worker/replica:0/task:0/cpu:0"
    worker1_device = "/job:worker/replica:0/task:1/cpu:0"

    @eager_def_function.function
    def f(a, b):
      return a + b

    with session.Session(workers[0].target) as sess:
      with ops.device(worker0_device):
        a = variable_scope.get_variable(
            "a", initializer=constant_op.constant(1.), use_resource=True)
      with ops.device(worker1_device):
        b = variable_scope.get_variable(
            "b", initializer=constant_op.constant(1.), use_resource=True)

      sess.run(variables.global_variables_initializer())

    config = config_pb2.ConfigProto()
    config.experimental.share_cluster_devices_in_session = True

    with session.Session(workers[0].target, config=config) as sess:
      res = sess.run(f(a, b))

    self.assertEqual(res, 2)

  def testCaptureHashTableInSharedIterator(self):
    worker, _ = test_util.create_local_cluster(1, 1)

    # NOTE(mrry): We must use the V2 variants of `HashTable`
    # etc. because these produce a `tf.resource`-typed output that is
    # compatible with the in-graph function implementation.
    default_val = -1
    keys = constant_op.constant(["brain", "salad", "surgery"])
    values = constant_op.constant([0, 1, 2], dtypes.int64)
    table = lookup_ops.HashTable(
        lookup_ops.KeyValueTensorInitializer(keys, values),
        default_val,
        shared_name="shared_table")

    input_sentences = dataset_ops.Dataset.from_tensor_slices(
        ["brain brain tank salad surgery", "surgery brain"])

    iterator = (
        input_sentences.map(lambda x: string_ops.string_split([x]).values).map(
            table.lookup)
        .make_initializable_iterator(shared_name="shared_iterator"))
    init_op = iterator.initializer
    get_next = iterator.get_next()

    with session.Session(worker[0].target) as sess:
      self.evaluate(table.initializer)
      self.evaluate(init_op)
      self.assertAllEqual([0, 0, -1, 1, 2], self.evaluate(get_next))

    with session.Session(worker[0].target) as sess:
      self.assertAllEqual([2, 0], self.evaluate(get_next))
      with self.assertRaises(errors.OutOfRangeError):
        sess.run(get_next)

  def testRemoteFunction(self):
    worker_config = config_pb2.ConfigProto()
    worker_config.device_count["CPU"] = 2
    worker, _ = test_util.create_local_cluster(
        1, 1, worker_config=worker_config)

    @function.Defun(dtypes.int32, dtypes.int32)
    def _remote_fn(a, b):
      return math_ops.multiply(a, b)

    with ops.device("/job:ps/task:0"):
      a = variables.Variable(2, dtype=dtypes.int32)
      b = variables.Variable(3, dtype=dtypes.int32)

    with ops.device("/job:worker/replica:0/task:0/cpu:0"):
      remote_op = functional_ops.remote_call(
          args=[a, b],
          Tout=[dtypes.int32],
          f=_remote_fn,
          target="/job:worker/replica:0/task:0/cpu:1")

    with session.Session(worker[0].target) as sess:
      self.evaluate(variables.global_variables_initializer())
      mul = self.evaluate(remote_op)
      self.assertEqual(mul, [6])

  def testRemoteIteratorUsingRemoteCallOp(self):
    worker_config = config_pb2.ConfigProto()
    worker_config.device_count["CPU"] = 2
    worker, _ = test_util.create_local_cluster(
        1, 1, worker_config=worker_config)

    self._testRemoteIteratorHelper("/job:worker/replica:0/task:0/cpu:0",
                                   "/job:worker/replica:0/task:0/cpu:1",
                                   worker[0].target)

 def setUpClass(cls):
   # We have to create a global in-process cluster because once an in-process
   # tensorflow server is created, there is no way to terminate it. Please see
   # multi_worker_test_base.py for more details.
   cls._workers, cls._ps = test_util.create_local_cluster(
       NUM_WORKERS, num_ps=NUM_PS)
   cls._cluster_spec = {
       WORKER: [_bytes_to_str(w.target) for w in cls._workers],
       PS: [_bytes_to_str(ps.target) for ps in cls._ps]
   }

 def _test_device_and_input_device_are_colocated(self, strategy):
   if context.executing_eagerly():
     self.skipTest(
         "cross-device tests are not supported with eager execution.")
   workers, _ = test_util.create_local_cluster(2, 0)
   inputs = strategy.make_input_fn_iterator(
       lambda _: dataset_ops.Dataset.range(5))
   comm_fn = lambda x: x + 1
   run_op = strategy.experimental_run(comm_fn, inputs)
   with session_lib.Session(target=workers[1].target) as sess:
     sess.run(inputs.initialize())
     sess.run(run_op)

 def testSerialize(self):
   worker = test_util.create_local_cluster(num_workers=1, num_ps=1)[0][0]
   with ops.Graph().as_default(), session.Session(target=worker.target):
     with ops.device("/job:worker"):
       t = constant_op.constant([[1.0], [2.0]])
       l = list_ops.tensor_list_from_tensor(t, element_shape=[1])
     with ops.device("/job:ps"):
       l_ps = array_ops.identity(l)
       l_ps, e = list_ops.tensor_list_pop_back(
           l_ps, element_dtype=dtypes.float32)
     with ops.device("/job:worker"):
       worker_e = array_ops.identity(e)
     self.assertAllEqual(self.evaluate(worker_e), [2.0])

 def testSerializeListWithUnknownRank(self):
   worker = test_util.create_local_cluster(num_workers=1, num_ps=1)[0][0]
   with ops.Graph().as_default(), session.Session(target=worker.target):
     with ops.device("/job:worker"):
       t = constant_op.constant([[1.0], [2.0]])
       l = list_ops.tensor_list_from_tensor(t, element_shape=None)
     with ops.device("/job:ps"):
       l_ps = array_ops.identity(l)
       element_shape = list_ops.tensor_list_element_shape(
           l_ps, shape_type=dtypes.int32)
     with ops.device("/job:worker"):
       element_shape = array_ops.identity(element_shape)
     self.assertEqual(self.evaluate(element_shape), -1)

  def testRemoteIteratorUsingRemoteCallOp(self):
    worker_config = config_pb2.ConfigProto()
    worker_config.device_count["CPU"] = 2
    worker, _ = test_util.create_local_cluster(
        1, 1, worker_config=worker_config)

    with ops.device("/job:worker/replica:0/task:0/cpu:1"):
      dataset_3 = dataset_ops.Dataset.from_tensor_slices([1, 2, 3])
      iterator_3 = dataset_3.make_one_shot_iterator()
      iterator_3_handle = iterator_3.string_handle()

    @function.Defun(dtypes.string)
    def _remote_fn(h):
      remote_iterator = dataset_ops.Iterator.from_string_handle(
          h, dataset_3.output_types, dataset_3.output_shapes)
      return remote_iterator.get_next()

    with ops.device("/job:worker/replica:0/task:0/cpu:0"):
      target_placeholder = array_ops.placeholder(dtypes.string, shape=[])
      remote_op = functional_ops.remote_call(
          args=[iterator_3_handle],
          Tout=[dtypes.int32],
          f=_remote_fn,
          target=target_placeholder)

    with session.Session(worker[0].target) as sess:
      elem = sess.run(
          remote_op,
          feed_dict={target_placeholder: "/job:worker/replica:0/task:0/cpu:1"})
      self.assertEqual(elem, [1])
      # Fails when target is cpu:0 where the resource is not located.
      with self.assertRaises(errors.InvalidArgumentError):
        sess.run(
            remote_op,
            feed_dict={
                target_placeholder: "/job:worker/replica:0/task:0/cpu:0"
            })
      elem = sess.run(
          remote_op,
          feed_dict={target_placeholder: "/job:worker/replica:0/task:0/cpu:1"})
      self.assertEqual(elem, [2])
      elem = sess.run(
          remote_op,
          feed_dict={target_placeholder: "/job:worker/replica:0/task:0/cpu:1"})
      self.assertEqual(elem, [3])
      with self.assertRaises(errors.OutOfRangeError):
        sess.run(
            remote_op,
            feed_dict={
                target_placeholder: "/job:worker/replica:0/task:0/cpu:1"
            })

  def setUpClass(cls):
    """Create a local cluster with 2 workers."""
    num_workers = 2
    # Leave some memory for cuda runtime.
    gpu_mem_frac = 0.7 / num_workers
    default_config = config_pb2.ConfigProto()
    default_config.gpu_options.per_process_gpu_memory_fraction = gpu_mem_frac

    # The local cluster takes some portion of the local GPUs and there is no way
    # for the cluster to terminate unless using multiple processes. Therefore,
    # we have to only create only one cluster throughout a test process.
    workers, _ = test_util.create_local_cluster(
        num_workers, num_ps=0, worker_config=default_config)
    cls._master_target = workers[0].target

 def testSerializeListWithInvalidTensors(self):
   worker = test_util.create_local_cluster(num_workers=1, num_ps=1)[0][0]
   with ops.Graph().as_default(), session.Session(target=worker.target):
     with ops.device("/job:worker"):
       l = list_ops.tensor_list_reserve(
           element_dtype=dtypes.float32, element_shape=[], num_elements=2)
       l = list_ops.tensor_list_set_item(l, 0, 1.)
     with ops.device("/job:ps"):
       l_ps = array_ops.identity(l)
       l_ps = list_ops.tensor_list_set_item(l_ps, 1, 2.)
       t = list_ops.tensor_list_stack(l_ps, element_dtype=dtypes.float32)
     with ops.device("/job:worker"):
       worker_t = array_ops.identity(t)
     self.assertAllEqual(self.evaluate(worker_t), [1.0, 2.0])

 def _test_device_and_input_device_are_colocated_with_function(self, strategy):
   if context.executing_eagerly():
     self.skipTest(
         "cross-device tests are not supported with eager execution.")
   workers, _ = test_util.create_local_cluster(2, 0)
   inputs = strategy.make_input_fn_iterator(
       lambda _: dataset_ops.Dataset.range(5))
   comm_fn = lambda x: x + 1
   experimental_run = def_function.function()(strategy.experimental_run)
   with ops.device("/job:worker/replica:0/task:1/device:CPU:0"):
     # The tf.function must be defined on the right device as well.
     run_op = experimental_run(comm_fn, inputs)
   with session_lib.Session(target=workers[1].target) as sess:
     sess.run(inputs.initialize())
     sess.run(run_op)

 def testSerializeListWithMaxNumElements(self):
   if context.num_gpus():
     # TODO(b/119151861): Enable on GPU.
     return
   worker = test_util.create_local_cluster(num_workers=1, num_ps=1)[0][0]
   with ops.Graph().as_default(), session.Session(target=worker.target):
     with ops.device("/job:worker"):
       l = list_ops.empty_tensor_list(
           element_shape=-1, element_dtype=dtypes.float32, max_num_elements=2)
       l = list_ops.tensor_list_push_back(l, 1.)
     with ops.device("/job:ps"):
       l_ps = array_ops.identity(l)
       l_ps = list_ops.tensor_list_push_back(l_ps, 2.)
     with self.assertRaisesRegexp(errors.InvalidArgumentError,
                                  "Tried to push item into a full list"):
       with ops.device("/job:worker"):
         l_worker = array_ops.identity(l_ps)
         l_worker = list_ops.tensor_list_push_back(l_worker, 3.0)
         self.evaluate(l_worker)