def test_get_size_info(self, num_clients):
@computations.federated_computation(
type_factory.at_clients(computation_types.SequenceType(tf.float32)),
type_factory.at_server(tf.float32))
def comp(temperatures, threshold):
client_data = [temperatures, intrinsics.federated_broadcast(threshold)]
result_map = intrinsics.federated_map(
count_over, intrinsics.federated_zip(client_data))
count_map = intrinsics.federated_map(count_total, temperatures)
return intrinsics.federated_mean(result_map, count_map)
factory = executor_stacks.sizing_executor_factory(num_clients=num_clients)
default_executor.set_default_executor(factory)
to_float = lambda x: tf.cast(x, tf.float32)
temperatures = [tf.data.Dataset.range(10).map(to_float)] * num_clients
threshold = 15.0
comp(temperatures, threshold)
# Each client receives a tf.float32 and uploads two tf.float32 values.
expected_broadcast_bits = num_clients * 32
expected_aggregate_bits = expected_broadcast_bits * 2
expected = ({
(('CLIENTS', num_clients),): [[1, tf.float32]] * num_clients
}, {
(('CLIENTS', num_clients),): [[1, tf.float32]] * num_clients * 2
}, [expected_broadcast_bits], [expected_aggregate_bits])
self.assertEqual(expected, factory.get_size_info())
def create_sizing_execution_context(default_num_clients: int = 0,
max_fanout: int = 100,
clients_per_thread: int = 1):
"""Creates an execution context that executes computations locally."""
factory = executor_stacks.sizing_executor_factory(
default_num_clients=default_num_clients,
max_fanout=max_fanout,
clients_per_thread=clients_per_thread)
return execution_context.ExecutionContext(
executor_fn=factory, compiler_fn=compiler.transform_to_native_form)
def _create_concurrent_maxthread_tuples():
tuples = []
for concurrency in range(1, 5):
local_ex_string = 'local_executor_{}_client_thread'.format(concurrency)
ex_factory = executor_stacks.local_executor_factory(
num_client_executors=concurrency)
tuples.append((local_ex_string, ex_factory, concurrency))
sizing_ex_string = 'sizing_executor_{}_client_thread'.format(concurrency)
ex_factory = executor_stacks.sizing_executor_factory(
num_client_executors=concurrency)
tuples.append((sizing_ex_string, ex_factory, concurrency))
return tuples
class ExecutorsTest(parameterized.TestCase):
@executor_test_utils.executors
def test_without_arguments(self):
@computations.tf_computation(tf.int32)
def add_one(x):
return x + 1
result = add_one(5)
self.assertEqual(result, 6)
@executor_test_utils.executors()
def test_with_no_arguments(self):
@computations.tf_computation(tf.int32)
def add_one(x):
return x + 1
result = add_one(5)
self.assertEqual(result, 6)
@executor_test_utils.executors(
('local', executor_stacks.local_executor_factory()),)
def test_with_one_argument(self):
@computations.tf_computation(tf.int32)
def add_one(x):
return x + 1
result = add_one(5)
self.assertEqual(result, 6)
@executor_test_utils.executors(
('local', executor_stacks.local_executor_factory()),
('sizing', executor_stacks.sizing_executor_factory()),
)
def test_with_two_argument(self):
@computations.tf_computation(tf.int32)
def add_one(x):
return x + 1
result = add_one(5)
self.assertEqual(result, 6)
def decorator(fn, *named_executors):
if not named_executors:
named_executors = [
('local', executor_stacks.local_executor_factory()),
('sizing', executor_stacks.sizing_executor_factory()),
]
@parameterized.named_parameters(*named_executors)
def wrapped_fn(self, executor):
"""Install a particular execution context before running `fn`."""
context = execution_context.ExecutionContext(executor)
with context_stack_impl.context_stack.install(context):
fn(self)
return wrapped_fn
def test_get_size_info(self, num_clients):
@computations.federated_computation(
type_factory.at_clients(computation_types.SequenceType(
tf.float32)), type_factory.at_server(tf.float32))
def comp(temperatures, threshold):
client_data = [
temperatures,
intrinsics.federated_broadcast(threshold)
]
result_map = intrinsics.federated_map(
count_over, intrinsics.federated_zip(client_data))
count_map = intrinsics.federated_map(count_total, temperatures)
return intrinsics.federated_mean(result_map, count_map)
sizing_factory = executor_stacks.sizing_executor_factory(
num_clients=num_clients)
sizing_context = execution_context.ExecutionContext(sizing_factory)
with get_context_stack.get_context_stack().install(sizing_context):
to_float = lambda x: tf.cast(x, tf.float32)
temperatures = [tf.data.Dataset.range(10).map(to_float)
] * num_clients
threshold = 15.0
comp(temperatures, threshold)
# Each client receives a tf.float32 and uploads two tf.float32 values.
expected_broadcast_bits = [num_clients * 32]
expected_aggregate_bits = [num_clients * 32 * 2]
expected_broadcast_history = {
(('CLIENTS', num_clients), ): [[1, tf.float32]] * num_clients
}
expected_aggregate_history = {
(('CLIENTS', num_clients), ):
[[1, tf.float32]] * num_clients * 2
}
size_info = sizing_factory.get_size_info()
self.assertEqual(expected_broadcast_history,
size_info.broadcast_history)
self.assertEqual(expected_aggregate_history,
size_info.aggregate_history)
self.assertEqual(expected_broadcast_bits, size_info.broadcast_bits)
self.assertEqual(expected_aggregate_bits, size_info.aggregate_bits)
def decorator(fn, *named_executors):
if isinstance(fn, type):
raise TypeError('Do not directly decorate classes with the executors '
'decorator; this will cause the tests to be skipped. '
'Decorate the member test functions instead.')
if not named_executors:
named_executors = [
('local', executor_stacks.local_executor_factory()),
('sizing', executor_stacks.sizing_executor_factory()),
]
@parameterized.named_parameters(*named_executors)
def wrapped_fn(self, executor):
"""Install a particular execution context before running `fn`."""
context = execution_context.ExecutionContext(executor)
with context_stack_impl.context_stack.install(context):
fn(self)
return wrapped_fn
def _create_concurrent_maxthread_tuples():
tuples = []
for concurrency in range(1, 5):
local_ex_string = 'local_executor_{}_clients_per_thread'.format(
concurrency)
tf_executor_mock = ExecutorMock()
ex_factory = executor_stacks.local_executor_factory(
clients_per_thread=concurrency, leaf_executor_fn=tf_executor_mock)
tuples.append(
(local_ex_string, ex_factory, concurrency, tf_executor_mock))
sizing_ex_string = 'sizing_executor_{}_client_thread'.format(
concurrency)
tf_executor_mock = ExecutorMock()
ex_factory = executor_stacks.sizing_executor_factory(
clients_per_thread=concurrency, leaf_executor_fn=tf_executor_mock)
tuples.append(
(sizing_ex_string, ex_factory, concurrency, tf_executor_mock))
debug_ex_string = 'debug_executor_{}_client_thread'.format(concurrency)
tf_executor_mock = ExecutorMock()
ex_factory = executor_stacks.thread_debugging_executor_factory(
clients_per_thread=concurrency, leaf_executor_fn=tf_executor_mock)
tuples.append(
(debug_ex_string, ex_factory, concurrency, tf_executor_mock))
return tuples
class ExecutorStacksTest(parameterized.TestCase):
@parameterized.named_parameters(
('local_executor', executor_stacks.local_executor_factory),
('sizing_executor', executor_stacks.sizing_executor_factory),
('debug_executor', executor_stacks.thread_debugging_executor_factory),
)
def test_construction_with_no_args(self, executor_factory_fn):
executor_factory_impl = executor_factory_fn()
self.assertIsInstance(executor_factory_impl,
executor_stacks.ResourceManagingExecutorFactory)
@parameterized.named_parameters(
('local_executor', executor_stacks.local_executor_factory),
('sizing_executor', executor_stacks.sizing_executor_factory),
)
def test_construction_raises_with_max_fanout_one(self,
executor_factory_fn):
with self.assertRaises(ValueError):
executor_factory_fn(max_fanout=1)
@parameterized.named_parameters(
('local_executor_none_clients',
executor_stacks.local_executor_factory()),
('sizing_executor_none_clients',
executor_stacks.sizing_executor_factory()),
('local_executor_three_clients',
executor_stacks.local_executor_factory(num_clients=3)),
('sizing_executor_three_clients',
executor_stacks.sizing_executor_factory(num_clients=3)),
)
@test_utils.skip_test_for_multi_gpu
def test_execution_of_temperature_sensor_example(self, executor):
comp = _temperature_sensor_example_next_fn()
to_float = lambda x: tf.cast(x, tf.float32)
temperatures = [
tf.data.Dataset.range(10).map(to_float),
tf.data.Dataset.range(20).map(to_float),
tf.data.Dataset.range(30).map(to_float),
]
threshold = 15.0
with executor_test_utils.install_executor(executor):
result = comp(temperatures, threshold)
self.assertAlmostEqual(result, 8.333, places=3)
@parameterized.named_parameters(
('local_executor', executor_stacks.local_executor_factory),
('sizing_executor', executor_stacks.sizing_executor_factory),
)
def test_execution_with_inferred_clients_larger_than_fanout(
self, executor_factory_fn):
@computations.federated_computation(
computation_types.at_clients(tf.int32))
def foo(x):
return intrinsics.federated_sum(x)
executor = executor_factory_fn(max_fanout=3)
with executor_test_utils.install_executor(executor):
result = foo([1, 2, 3, 4, 5, 6, 7, 8, 9, 10])
self.assertEqual(result, 55)
@parameterized.named_parameters(
('local_executor_none_clients',
executor_stacks.local_executor_factory()),
('sizing_executor_none_clients',
executor_stacks.sizing_executor_factory()),
('debug_executor_none_clients',
executor_stacks.thread_debugging_executor_factory()),
('local_executor_one_client',
executor_stacks.local_executor_factory(num_clients=1)),
('sizing_executor_one_client',
executor_stacks.sizing_executor_factory(num_clients=1)),
('debug_executor_one_client',
executor_stacks.thread_debugging_executor_factory(num_clients=1)),
)
def test_execution_of_tensorflow(self, executor):
@computations.tf_computation
def comp():
return tf.math.add(5, 5)
with executor_test_utils.install_executor(executor):
result = comp()
self.assertEqual(result, 10)
@parameterized.named_parameters(*_create_concurrent_maxthread_tuples())
def test_limiting_concurrency_constructs_one_eager_executor(
self, ex_factory, clients_per_thread, tf_executor_mock):
num_clients = 10
ex_factory.create_executor({placements.CLIENTS: num_clients})
concurrency_level = math.ceil(num_clients / clients_per_thread)
args_list = tf_executor_mock.call_args_list
# One for server executor, one for unplaced executor, concurrency_level for
# clients.
self.assertLen(args_list, concurrency_level + 2)
@mock.patch(
'tensorflow_federated.python.core.impl.executors.reference_resolving_executor.ReferenceResolvingExecutor',
return_value=ExecutorMock())
def test_thread_debugging_executor_constructs_exactly_one_reference_resolving_executor(
self, executor_mock):
executor_stacks.thread_debugging_executor_factory().create_executor(
{placements.CLIENTS: 10})
executor_mock.assert_called_once()
@parameterized.named_parameters(
('local_executor', executor_stacks.local_executor_factory),
('sizing_executor', executor_stacks.sizing_executor_factory),
('debug_executor', executor_stacks.thread_debugging_executor_factory),
)
def test_create_executor_raises_with_wrong_cardinalities(
self, executor_factory_fn):
executor_factory_impl = executor_factory_fn(num_clients=5)
cardinalities = {
placements.SERVER: 1,
None: 1,
placements.CLIENTS: 1,
}
with self.assertRaises(ValueError, ):
executor_factory_impl.create_executor(cardinalities)
class ExecutorStacksTest(parameterized.TestCase):
@parameterized.named_parameters(
('local_executor', executor_stacks.local_executor_factory),
('sizing_executor', executor_stacks.sizing_executor_factory),
)
def test_construction_with_no_args(self, executor_factory_fn):
executor_factory_impl = executor_factory_fn()
self.assertIsInstance(executor_factory_impl,
executor_factory.ExecutorFactoryImpl)
@parameterized.named_parameters(
('local_executor', executor_stacks.local_executor_factory),
('sizing_executor', executor_stacks.sizing_executor_factory),
)
def test_construction_raises_with_max_fanout_one(self,
executor_factory_fn):
with self.assertRaises(ValueError):
executor_factory_fn(max_fanout=1)
@parameterized.named_parameters(
('local_executor_none_clients',
executor_stacks.local_executor_factory()),
('sizing_executor_none_clients',
executor_stacks.sizing_executor_factory()),
('local_executor_three_clients',
executor_stacks.local_executor_factory(num_clients=3)),
('sizing_executor_three_clients',
executor_stacks.sizing_executor_factory(num_clients=3)),
)
def test_execution_of_temperature_sensor_example(self, executor):
comp = _temperature_sensor_example_next_fn()
to_float = lambda x: tf.cast(x, tf.float32)
temperatures = [
tf.data.Dataset.range(10).map(to_float),
tf.data.Dataset.range(20).map(to_float),
tf.data.Dataset.range(30).map(to_float),
]
threshold = 15.0
with executor_test_utils.install_executor(executor):
result = comp(temperatures, threshold)
self.assertAlmostEqual(result, 8.333, places=3)
@parameterized.named_parameters(
('local_executor', executor_stacks.local_executor_factory),
('sizing_executor', executor_stacks.sizing_executor_factory),
)
def test_execution_with_inferred_clients_larger_than_fanout(
self, executor_factory_fn):
@computations.federated_computation(type_factory.at_clients(tf.int32))
def foo(x):
return intrinsics.federated_sum(x)
executor = executor_factory_fn(max_fanout=3)
with executor_test_utils.install_executor(executor):
result = foo([1, 2, 3, 4, 5, 6, 7, 8, 9, 10])
self.assertEqual(result, 55)
@parameterized.named_parameters(
('local_executor_none_clients',
executor_stacks.local_executor_factory()),
('sizing_executor_none_clients',
executor_stacks.sizing_executor_factory()),
('local_executor_one_client',
executor_stacks.local_executor_factory(num_clients=1)),
('sizing_executor_one_client',
executor_stacks.sizing_executor_factory(num_clients=1)),
)
def test_execution_of_tensorflow(self, executor):
@computations.tf_computation
def comp():
return tf.math.add(5, 5)
with executor_test_utils.install_executor(executor):
result = comp()
self.assertEqual(result, 10)
@parameterized.named_parameters(
('local_executor', executor_stacks.local_executor_factory),
('sizing_executor', executor_stacks.sizing_executor_factory),
)
def test_create_executor_raises_with_wrong_cardinalities(
self, executor_factory_fn):
executor_factory_impl = executor_factory_fn(num_clients=5)
cardinalities = {
placement_literals.SERVER: 1,
None: 1,
placement_literals.CLIENTS: 1,
}
with self.assertRaises(ValueError, ):
executor_factory_impl.create_executor(cardinalities)
