def bufferize(worker: AbstractWorker, state: "State") -> StatePB: """ Serialize the State to Protobuf message """ protobuf_state = StatePB() protobuf_placeholders = [ sy.serde.protobuf.serde._bufferize(worker, placeholder) for placeholder in state.state_placeholders ] protobuf_state.placeholders.extend(protobuf_placeholders) state_tensors = [] for tensor in state.tensors(): protobuf_tensor = sy.serde.protobuf.serde._bufferize( worker, tensor) state_tensor = StateTensorPB() if type(protobuf_tensor) == ParameterPB: state_tensor.torch_param.CopyFrom( sy.serde.protobuf.serde._bufferize(worker, tensor)) else: state_tensor.torch_tensor.CopyFrom( sy.serde.protobuf.serde._bufferize(worker, tensor)) state_tensors.append(state_tensor) protobuf_state.tensors.extend(state_tensors) return protobuf_state
def _object2proto(self) -> StatePB: """Returns a protobuf serialization of self. As a requirement of all objects which inherit from Serializable, this method transforms the current object into the corresponding Protobuf object so that it can be further serialized. :return: returns a protobuf object :rtype: ObjectWithID_PB .. note:: This method is purely an internal method. Please use object.serialize() or one of the other public serialization methods if you wish to serialize an object. """ proto = StatePB() protobuf_placeholders = [ serialize(placeholder) for placeholder in self.state_placeholders ] proto.placeholders.extend(protobuf_placeholders) state_tensors = [] for tensor in self.tensors(): state_tensor = StateTensorPB() state_tensor.torch_tensor.CopyFrom(serialize_tensor(tensor)) state_tensors.append(state_tensor) proto.tensors.extend(state_tensors) return proto
def unserialize_model_params(bin: bin): """Unserializes model or checkpoint or diff stored in db to list of tensors""" state = StatePB() state.ParseFromString(bin) worker = sy.VirtualWorker(hook=None) state = protobuf.serde._unbufferize(worker, state) model_params = state.tensors() return model_params
async def main(): #auth_token = jwt.encode({}, private_key, algorithm='RS256').decode('ascii') #print(auth_token) auth_request = { "type": "model-centric/authenticate", "data": { "model_name": name, "model_version": version, #"auth_token": auth_token, } } print('Auth request: ', json.dumps(auth_request, indent=2)) auth_response = await sendWsMessage(auth_request) print('Auth response: ', json.dumps(auth_response, indent=2)) cycle_request = { "type": "model-centric/cycle-request", "data": { "worker_id": auth_response['data']['worker_id'], "model": name, "version": version, "ping": 1, "download": 10000, "upload": 10000, } } cycle_response = await sendWsMessage(cycle_request) print('Cycle response:', json.dumps(cycle_response, indent=2)) worker_id = auth_response['data']['worker_id'] request_key = cycle_response['data']['request_key'] model_id = cycle_response['data']['model_id'] training_plan_id = cycle_response['data']['plans']['training_plan'] # Model req = requests.get( f"http://{gridAddress}/model-centric/get-model?worker_id={worker_id}&request_key={request_key}&model_id={model_id}" ) model_data = req.content pb = StatePB() pb.ParseFromString(req.content) model_params_downloaded = protobuf.serde._unbufferize( hook.local_worker, pb) print("Params shapes:", [p.shape for p in model_params_downloaded.tensors()]) # Plan "list of ops" req = requests.get( f"http://{gridAddress}/model-centric/get-plan?worker_id={worker_id}&request_key={request_key}&plan_id={training_plan_id}&receive_operations_as=list" ) pb = PlanPB() pb.ParseFromString(req.content) plan_ops = protobuf.serde._unbufferize(hook.local_worker, pb) print(plan_ops.code) print(plan_ops.torchscript) # Plan "torchscript" req = requests.get( f"http://{gridAddress}/model-centric/get-plan?worker_id={worker_id}&request_key={request_key}&plan_id={training_plan_id}&receive_operations_as=torchscript" ) pb = PlanPB() pb.ParseFromString(req.content) plan_ts = protobuf.serde._unbufferize(hook.local_worker, pb) print(plan_ts.code) print(plan_ts.torchscript.code) # Plan "tfjs" req = requests.get( f"http://{gridAddress}/model-centric/get-plan?worker_id={worker_id}&request_key={request_key}&plan_id={training_plan_id}&receive_operations_as=tfjs" ) pb = PlanPB() pb.ParseFromString(req.content) plan_tfjs = protobuf.serde._unbufferize(hook.local_worker, pb) print(plan_tfjs.code)
}, } cycle_response = sendWsMessage(cycle_request) print("Cycle response:", json.dumps(cycle_response, indent=2)) worker_id = auth_response["data"]["worker_id"] request_key = cycle_response["data"]["request_key"] model_id = cycle_response["data"]["model_id"] training_plan_id = cycle_response["data"]["plans"]["training_plan"] # Let's download Model and Training Plan (in various trainslations) and check they are actually workable. req = requests.get( f"http://{config.GRID_ADDRESS}/model-centric/get-model?worker_id={worker_id}&request_key={request_key}&model_id={model_id}" ) model_data = req.content pb = StatePB() pb.ParseFromString(req.content) model_params_downloaded = protobuf.serde._unbufferize(hook.local_worker, pb) print("Params shapes:", [p.shape for p in model_params_downloaded.tensors()]) # Plan "list of ops" req = requests.get( f"http://{config.GRID_ADDRESS}/model-centric/get-plan?worker_id={worker_id}&request_key={request_key}&plan_id={training_plan_id}&receive_operations_as=list" ) pb = PlanPB() pb.ParseFromString(req.content) plan_ops = protobuf.serde._unbufferize(hook.local_worker, pb) print(plan_ops.code) print(plan_ops.torchscript) # Plan "torchscript"