def test_channel(self):
channel = make_insecure_channel('worker_0', mode=ChannelType.INTERNAL)
stub = tm_grpc.TrainerMasterServiceStub(channel)
request = tm_pb.DataBlockRequest()
print(stub.RequestDataBlock(request))
self.assertIsNotNone(stub.RequestDataBlock(request))
channel = make_insecure_channel('worker_0', mode=ChannelType.REMOTE)
stub = tm_grpc.TrainerMasterServiceStub(channel)
request = tm_pb.DataBlockRequest()
print(stub.RequestDataBlock(request))
self.assertIsNotNone(stub.RequestDataBlock(request))
def __init__(self, addr, role, task_id):
self._addr = addr
self._role = role
self._task_id = task_id
channel = make_insecure_channel(self._addr, ChannelType.INTERNAL)
self._stub = tm_grpc.TrainerMasterServiceStub(channel)
self._request = tm_pb.DataBlockRequest()
if self._role == 'leader':
self._request.worker_rank = self._task_id
def request_data_block(self, block_id):
request = tm_pb.DataBlockRequest(worker_rank=self._worker_rank,
block_id=block_id)
response = _grpc_with_retry(
lambda: self._client.RequestDataBlock(request))
if response.status.code == common_pb.StatusCode.STATUS_SUCCESS:
fl_logging.debug(
"succeeded to get datablock, id:%s, data_path: %s",
response.block_id, response.data_path)
return response
if response.status.code == \
common_pb.StatusCode.STATUS_INVALID_DATA_BLOCK:
fl_logging.error("invalid data block id: %s", request.block_id)
return None
if response.status.code == common_pb.StatusCode.STATUS_DATA_FINISHED:
fl_logging.info("data block finished")
return None
raise RuntimeError("RequestDataBlock error, code: %s, msg: %s"% \
(common_pb.StatusCode.Name(response.status.code),
response.status.error_message))