Skip to content
New issue

Have a question about this project? # for a free GitHub account to open an issue and contact its maintainers and the community.

#

By clicking “#”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? # to your account

send local task results to retry if connection is lost #465

Merged
merged 4 commits into from
Aug 12, 2022
Merged

send local task results to retry if connection is lost #465

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
7755f06
send local task results to retry if connection is lost
Jul 28, 2022
7dc7d87
remove unused imports
Jul 28, 2022
b9ef3c7
decorator for resending data request
Jul 29, 2022
d1292ce
add message to resume training at server
Aug 12, 2022
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
17 changes: 17 additions & 0 deletions openfl/transport/grpc/aggregator_client.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -85,6 +85,19 @@ def wrapper(self, *args, **kwargs):

return wrapper

def _resend_data_on_reconnection(func):
def wrapper(self, *args, **kwargs):
while True:
try:
response = func(self, *args, **kwargs)
except grpc.RpcError as e:
if e.code() == grpc.StatusCode.UNKNOWN:
self.logger.info(f'Attempting to resend data request to aggregator at {self.uri}')
continue
break
return response

return wrapper

class AggregatorGRPCClient:
"""Client to the aggregator over gRPC-TLS."""
Expand Down Expand Up @@ -258,6 +271,7 @@ def reconnect(self):
)

@_atomic_connection
@_resend_data_on_reconnection
def get_tasks(self, collaborator_name):
"""Get tasks from the aggregator."""
self._set_header(collaborator_name)
Expand All @@ -268,10 +282,12 @@ def get_tasks(self, collaborator_name):
return response.tasks, response.round_number, response.sleep_time, response.quit

@_atomic_connection
@_resend_data_on_reconnection
def get_aggregated_tensor(self, collaborator_name, tensor_name, round_number,
report, tags, require_lossless):
"""Get aggregated tensor from the aggregator."""
self._set_header(collaborator_name)

request = aggregator_pb2.GetAggregatedTensorRequest(
header=self.header,
tensor_name=tensor_name,
Expand All @@ -287,6 +303,7 @@ def get_aggregated_tensor(self, collaborator_name, tensor_name, round_number,
return response.tensor

@_atomic_connection
@_resend_data_on_reconnection
def send_local_task_results(self, collaborator_name, round_number,
task_name, data_size, named_tensors):
"""Send task results to the aggregator."""
Expand Down
9 changes: 6 additions & 3 deletions openfl/transport/grpc/aggregator_server.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -205,8 +205,11 @@ def SendLocalTaskResults(self, request, context): # NOQA:N802
context: The gRPC context

"""
proto = aggregator_pb2.TaskResults()
proto = utils.datastream_to_proto(proto, request)
try:
proto = aggregator_pb2.TaskResults()
proto = utils.datastream_to_proto(proto, request)
except RuntimeError:
raise RuntimeError('Empty stream message, reestablishing connection from client to resume training...')

self.validate_collaborator(proto, context)
# all messages get sanity checked
Expand All @@ -223,7 +226,7 @@ def SendLocalTaskResults(self, request, context): # NOQA:N802
return aggregator_pb2.SendLocalTaskResultsResponse(
header=self.get_header(collaborator_name)
)

def get_server(self):
"""Return gRPC server."""
self.server = server(ThreadPoolExecutor(max_workers=cpu_count()),
Expand Down