Adds failing integration test for fewer-clients-than-workers bug.
@@ -76,6 +76,40 @@ class RemoteRuntimeIntegrationTest(parameterized.TestCase):
result = map_add_one([0, 1])
self.assertEqual(result, [1, 2])
+ @parameterized.named_parameters((
+ 'native_remote',
+ remote_runtime_test_utils.create_localhost_remote_context(_WORKER_PORTS),
+ remote_runtime_test_utils.create_localhost_worker_contexts(_WORKER_PORTS),
+ list(range(len(_WORKER_PORTS) - 1)),
+ ), (
+ 'native_remote_intermediate_aggregator',
+ remote_runtime_test_utils.create_localhost_remote_context(
+ _AGGREGATOR_PORTS),
+ remote_runtime_test_utils.create_localhost_aggregator_contexts(
+ _WORKER_PORTS, _AGGREGATOR_PORTS),
+ list(range(len(_AGGREGATOR_PORTS) - 1)),
+ ))
+ def test_computations_run_with_fewer_clients_than_remote_connections(
+ self, context, serving_contexts, arg):
+ self.skipTest('b/170315887')
+
+ @tff.tf_computation(tf.int32)
+ def add_one(x):
+ return x + 1
+
+ @tff.federated_computation(tff.type_at_clients(tf.int32))
+ def map_add_one(federated_arg):
+ return tff.federated_map(add_one, federated_arg)
+
+ context_stack = tff.framework.get_context_stack()
+ with context_stack.install(context):
+
+ with contextlib.ExitStack() as stack:
+ for server_context in serving_contexts:
+ stack.enter_context(server_context)
+ result = map_add_one(arg)
+ self.assertEqual(result, [x + 1 for x in arg])
+
if __name__ == '__main__':
absltest.main()
