def test_monkey_patch_default_variable_placement_strategy(
self, num_tasks, op_names, before_want_ps, after_want_ps):
"""Checks that ps placement is based on var name."""
var_ops = [tf.Variable(0., name=op_name).op for op_name in op_names]
before_device_fn = tf.compat.v1.train.replica_device_setter(
ps_tasks=num_tasks)
self.assertEqual(before_want_ps, [before_device_fn(op) for op in var_ops])
with monkey_patch_default_variable_placement_strategy():
after_device_fn = tf.compat.v1.train.replica_device_setter(
ps_tasks=num_tasks)
self.assertEqual(after_want_ps, [after_device_fn(op) for op in var_ops])
# Check that monkey-patch is only for the context.
before_device_fn = tf.compat.v1.train.replica_device_setter(
ps_tasks=num_tasks)
self.assertEqual(before_want_ps, [before_device_fn(op) for op in var_ops])
def test_monkey_patch_default_variable_placement_strategy_no_ps(self):
with monkey_patch_default_variable_placement_strategy():
device_fn = tf.train.replica_device_setter(ps_tasks=0)
self.assertIsNone(device_fn)