Add distributed capability for FLEX

[mariogmarq]

This pull request provides FLEX with the distributed framework which allows to jointly train models with different devices. Code is self contained, adds gRPC as a dependency and provides a high coverage with integration and unit tests.

[github-actions]

Coverage Report

File	Stmts	Miss	Cover	Missing
flex/common
utils.py	75	15	80%	135–151, 189
flex/data
dataset.py	155	37	76%	53–59, 201, 203, 215, 223–239, 307–308, 315, 323, 328, 335–336, 354–371, 398, 414–415
dataset_pt_utils.py	29	2	93%	48, 59
fed_data_distribution.py	182	3	98%	50–51, 397
lazy_indexable.py	79	1	99%	89
pluggable_datasets.py	62	9	85%	27, 41–49
flex/datasets
federated_datasets.py	45	26	42%	35, 39–71, 75–89, 102
flexible_datasets.py	11	3	73%	29, 46–47
flex/distributed
__init__.py	3	2	33%	18–19
client.py	115	104	10%	30–238
client_builder.py	72	72	0%	18–252
common.py	16	13	19%	24–42
server.py	170	170	0%	17–487
flex/distributed/aio
__init__.py	1	1	0%	17
server.py	151	151	0%	17–479
flex/distributed/proto
tensor_pb2.py	13	9	31%	12–25
tensor_pb2_grpc.py	14	14	0%	3–20
transport_pb2.py	50	50	0%	6–64
transport_pb2_grpc.py	32	32	0%	3–96
flex/pool
aggregators.py	50	4	92%	55–56, 145–146
primitives_pt.py	63	50	21%	46, 61–63, 104–122, 158–170, 197–206, 235–244
primitives_tf.py	42	7	83%	74–84
TOTAL	1923	775	60%

Tests	Skipped	Failures	Errors	Time
193	4 💤	4 ❌	3 🔥	2m 24s ⏱️