async def compute_federated_aggregate(
self,
arg: FederatedResolvingStrategyValue) -> FederatedResolvingStrategyValue:
val_type, zero_type, accumulate_type, merge_type, report_type = (
executor_utils.parse_federated_aggregate_argument_types(
arg.type_signature))
del val_type, merge_type
py_typecheck.check_type(arg.internal_representation, structure.Struct)
py_typecheck.check_len(arg.internal_representation, 5)
val, zero, accumulate, merge, report = arg.internal_representation
# Discard `merge`. Since all aggregation happens on a single executor,
# there's no need for this additional layer.
del merge
# Re-wrap `zero` in a `FederatingResolvingStrategyValue` to ensure that it
# is an `ExecutorValue` rather than a `Struct` (since the internal
# representation can include embedded values, lists of embedded values
# (in the case of federated values), or `Struct`s.
zero = FederatedResolvingStrategyValue(zero, zero_type)
pre_report = await self.reduce(val, zero, accumulate, accumulate_type)
py_typecheck.check_type(pre_report.type_signature,
computation_types.FederatedType)
pre_report.type_signature.member.check_equivalent_to(report_type.parameter)
return await self.compute_federated_apply(
FederatedResolvingStrategyValue(
structure.Struct([(None, report),
(None, pre_report.internal_representation)]),
computation_types.StructType(
(report_type, pre_report.type_signature))))
async def compute_federated_aggregate(
self, arg: FederatedResolvingStrategyValue
) -> FederatedResolvingStrategyValue:
val_type, zero_type, accumulate_type, merge_type, report_type = (
executor_utils.parse_federated_aggregate_argument_types(
arg.type_signature))
del val_type, zero_type, merge_type
py_typecheck.check_type(arg.internal_representation, structure.Struct)
py_typecheck.check_len(arg.internal_representation, 5)
val, zero, accumulate, merge, report = arg.internal_representation
# Discard `merge`. Since all aggregation happens on a single executor,
# there's no need for this additional layer.
del merge
pre_report = await self.reduce(val, zero, accumulate, accumulate_type)
py_typecheck.check_type(pre_report.type_signature,
computation_types.FederatedType)
pre_report.type_signature.member.check_equivalent_to(
report_type.parameter)
return await self.compute_federated_apply(
FederatedResolvingStrategyValue(
structure.Struct([(None, report),
(None, pre_report.internal_representation)]),
computation_types.StructType(
(report_type, pre_report.type_signature))))
async def compute_federated_aggregate(
self, arg: FederatedComposingStrategyValue
) -> FederatedComposingStrategyValue:
value_type, zero_type, accumulate_type, merge_type, report_type = (
executor_utils.parse_federated_aggregate_argument_types(
arg.type_signature))
py_typecheck.check_type(arg.internal_representation, structure.Struct)
py_typecheck.check_len(arg.internal_representation, 5)
val = arg.internal_representation[0]
py_typecheck.check_type(val, list)
py_typecheck.check_len(val, len(self._target_executors))
identity_report, identity_report_type = tensorflow_computation_factory.create_identity(
zero_type)
aggr_type = computation_types.FunctionType(
computation_types.StructType([
value_type, zero_type, accumulate_type, merge_type,
identity_report_type
]), computation_types.at_server(zero_type))
aggr_comp = executor_utils.create_intrinsic_comp(
intrinsic_defs.FEDERATED_AGGREGATE, aggr_type)
zero = await (await
self._executor.create_selection(arg, index=1)).compute()
accumulate = arg.internal_representation[2]
merge = arg.internal_representation[3]
report = arg.internal_representation[4]
async def _child_fn(ex, v):
py_typecheck.check_type(v, executor_value_base.ExecutorValue)
arg_values = [
ex.create_value(zero, zero_type),
ex.create_value(accumulate, accumulate_type),
ex.create_value(merge, merge_type),
ex.create_value(identity_report, identity_report_type)
]
aggr_func, aggr_args = await asyncio.gather(
ex.create_value(aggr_comp, aggr_type),
ex.create_struct([v] +
list(await asyncio.gather(*arg_values))))
child_result = await (await ex.create_call(aggr_func,
aggr_args)).compute()
result_at_server = await self._server_executor.create_value(
child_result, zero_type)
return result_at_server
val_futures = asyncio.as_completed(
[_child_fn(c, v) for c, v in zip(self._target_executors, val)])
parent_merge, parent_report = await asyncio.gather(
self._server_executor.create_value(merge, merge_type),
self._server_executor.create_value(report, report_type))
merge_result = await next(val_futures)
for next_val_future in val_futures:
next_val = await next_val_future
merge_arg = await self._server_executor.create_struct(
[merge_result, next_val])
merge_result = await self._server_executor.create_call(
parent_merge, merge_arg)
report_result = await self._server_executor.create_call(
parent_report, merge_result)
return FederatedComposingStrategyValue(
report_result, computation_types.at_server(report_type.result))
async def compute_federated_aggregate(
self, arg: FederatedComposingStrategyValue
) -> FederatedComposingStrategyValue:
value_type, zero_type, accumulate_type, merge_type, report_type = (
executor_utils.parse_federated_aggregate_argument_types(
arg.type_signature))
py_typecheck.check_type(arg.internal_representation,
anonymous_tuple.AnonymousTuple)
py_typecheck.check_len(arg.internal_representation, 5)
val = arg.internal_representation[0]
py_typecheck.check_type(val, list)
py_typecheck.check_len(val, len(self._target_executors))
identity_report = tensorflow_computation_factory.create_identity(
zero_type)
identity_report_type = type_factory.unary_op(zero_type)
aggr_type = computation_types.FunctionType(
computation_types.NamedTupleType([
value_type, zero_type, accumulate_type, merge_type,
identity_report_type
]), type_factory.at_server(zero_type))
aggr_comp = executor_utils.create_intrinsic_comp(
intrinsic_defs.FEDERATED_AGGREGATE, aggr_type)
zero = await (await
self._executor.create_selection(arg, index=1)).compute()
accumulate = arg.internal_representation[2]
merge = arg.internal_representation[3]
report = arg.internal_representation[4]
async def _child_fn(ex, v):
py_typecheck.check_type(v, executor_value_base.ExecutorValue)
aggr_func, aggr_args = await asyncio.gather(
ex.create_value(aggr_comp, aggr_type),
ex.create_tuple([v] + list(await asyncio.gather(
ex.create_value(zero, zero_type),
ex.create_value(accumulate, accumulate_type),
ex.create_value(merge, merge_type),
ex.create_value(identity_report, identity_report_type)))))
return await (await ex.create_call(aggr_func, aggr_args)).compute()
vals = await asyncio.gather(
*[_child_fn(c, v) for c, v in zip(self._target_executors, val)])
parent_vals = await asyncio.gather(
*[self._server_executor.create_value(v, zero_type) for v in vals])
parent_merge, parent_report = await asyncio.gather(
self._server_executor.create_value(merge, merge_type),
self._server_executor.create_value(report, report_type))
merge_result = parent_vals[0]
for next_val in parent_vals[1:]:
merge_result = await self._server_executor.create_call(
parent_merge, await
self._server_executor.create_tuple([merge_result, next_val]))
return FederatedComposingStrategyValue(
await self._server_executor.create_call(parent_report,
merge_result),
type_factory.at_server(report_type.result))
async def compute_federated_aggregate(
self, arg: FederatedResolvingStrategyValue
) -> FederatedResolvingStrategyValue:
val_type, zero_type, accumulate_type, _, report_type = (
executor_utils.parse_federated_aggregate_argument_types(
arg.type_signature))
py_typecheck.check_type(arg.internal_representation,
anonymous_tuple.AnonymousTuple)
py_typecheck.check_len(arg.internal_representation, 5)
# Note: This is a simple initial implementation that simply forwards this
# to `federated_reduce()`. The more complete implementation would be able
# to take advantage of the parallelism afforded by `merge` to reduce the
# cost from liner (with respect to the number of clients) to sub-linear.
# TODO(b/134543154): Expand this implementation to take advantage of the
# parallelism afforded by `merge`.
val = arg.internal_representation[0]
zero = arg.internal_representation[1]
accumulate = arg.internal_representation[2]
pre_report = await self.compute_federated_reduce(
FederatedResolvingStrategyValue(
anonymous_tuple.AnonymousTuple([(None, val), (None, zero),
(None, accumulate)]),
computation_types.NamedTupleType(
(val_type, zero_type, accumulate_type))))
py_typecheck.check_type(pre_report.type_signature,
computation_types.FederatedType)
pre_report.type_signature.member.check_equivalent_to(
report_type.parameter)
report = arg.internal_representation[4]
return await self.compute_federated_apply(
FederatedResolvingStrategyValue(
anonymous_tuple.AnonymousTuple([
(None, report), (None, pre_report.internal_representation)
]),
computation_types.NamedTupleType(
(report_type, pre_report.type_signature))))
