def __init__(self, adapter: RollSiteAdapter, options: dict = None): if options is None: options = {} self.adapter = adapter self.roll_site_header: ErRollSiteHeader = adapter.roll_site_header self.namespace = adapter.namespace self.name = create_store_name(self.roll_site_header) self.tagged_key = '' self.obj_type = adapter.obj_type self.proxy_endpoint = adapter.proxy_endpoint channel = self.grpc_channel_factory.create_channel(self.proxy_endpoint) self.stub = proxy_pb2_grpc.DataTransferServiceStub(channel) static_er_conf = get_static_er_conf() self.__bin_packet_len = int(options.get( RollSiteConfKeys.EGGROLL_ROLLSITE_ADAPTER_SENDBUF_SIZE.key, static_er_conf.get(RollSiteConfKeys.EGGROLL_ROLLSITE_ADAPTER_SENDBUF_SIZE.key, RollSiteConfKeys.EGGROLL_ROLLSITE_ADAPTER_SENDBUF_SIZE.default_value))) self.total_written = 0 self.ba = bytearray(self.__bin_packet_len) self.buffer = ArrayByteBuffer(self.ba) self.writer = PairBinWriter(pair_buffer=self.buffer) self.push_batch_cnt = 0 self.push_pair_cnt = 0 self.topic_src = proxy_pb2.Topic(name=self.name, partyId=self.roll_site_header._src_party_id, role=self.roll_site_header._src_role, callback=None) self.topic_dst = proxy_pb2.Topic(name=self.name, partyId=self.roll_site_header._dst_party_id, role=self.roll_site_header._dst_role, callback=None)
def init_job_session_pair(self, roll_site_session_id, er_session_id): try: task_info = proxy_pb2.Task(model=proxy_pb2.Model( name=roll_site_session_id, dataKey=bytes(er_session_id, encoding='utf8'))) topic_src = proxy_pb2.Topic(name="init_job_session_pair", partyId=self.party_id, role=self.role, callback=None) topic_dst = proxy_pb2.Topic(name="init_job_session_pair", partyId=self.party_id, role=self.role, callback=None) command_test = proxy_pb2.Command(name="init_job_session_pair") conf_test = proxy_pb2.Conf(overallTimeout=1000, completionWaitTimeout=1000, packetIntervalTimeout=1000, maxRetries=10) metadata = proxy_pb2.Metadata(task=task_info, src=topic_src, dst=topic_dst, command=command_test, operator="init_job_session_pair", seq=0, ack=0) packet = proxy_pb2.Packet(header=metadata) self.stub.unaryCall(packet) L.info( f"send RollSiteContext init to Proxy: {to_one_line_string(packet)}" ) except Exception as e: raise GrpcCallError("init_job_session_pair", self.proxy_endpoint, e)
def wrap_grpc_packet(_json_body, _method, _url, _src_party_id, _dst_party_id, job_id=None, overall_timeout=DEFAULT_GRPC_OVERALL_TIMEOUT): _src_end_point = basic_meta_pb2.Endpoint(ip=IP, port=GRPC_PORT) _src = proxy_pb2.Topic(name=job_id, partyId="{}".format(_src_party_id), role=ROLE, callback=_src_end_point) _dst = proxy_pb2.Topic(name=job_id, partyId="{}".format(_dst_party_id), role=ROLE, callback=None) _task = proxy_pb2.Task(taskId=job_id) _command = proxy_pb2.Command(name=ROLE) _conf = proxy_pb2.Conf(overallTimeout=overall_timeout) _meta = proxy_pb2.Metadata(src=_src, dst=_dst, task=_task, command=_command, operator=_method, conf=_conf) _data = proxy_pb2.Data(key=_url, value=bytes(json.dumps(_json_body), 'utf-8')) return proxy_pb2.Packet(header=_meta, body=_data)
def pull(self, parties: list = None): self._pull_start_wall_time = time.time() self._pull_start_cpu_time = time.perf_counter() futures = [] for src_role, src_party_id in parties: src_party_id = str(src_party_id) roll_site_header = ErRollSiteHeader( roll_site_session_id=self.roll_site_session_id, name=self.name, tag=self.tag, src_role=src_role, src_party_id=src_party_id, dst_role=self.local_role, dst_party_id=self.party_id) _tagged_key = create_store_name(roll_site_header) name = _tagged_key model = proxy_pb2.Model(name=_stringify(roll_site_header)) task_info = proxy_pb2.Task(taskId=name, model=model) topic_src = proxy_pb2.Topic(name="get_status", partyId=src_party_id, role=src_role, callback=None) topic_dst = proxy_pb2.Topic(name="get_status", partyId=self.party_id, role=self.local_role, callback=None) get_status_command = proxy_pb2.Command(name="get_status") conf_test = proxy_pb2.Conf(overallTimeout=1000, completionWaitTimeout=1000, packetIntervalTimeout=1000, maxRetries=10) metadata = proxy_pb2.Metadata(task=task_info, src=topic_src, dst=topic_dst, command=get_status_command, operator="getStatus", seq=0, ack=0) packet = proxy_pb2.Packet(header=metadata) namespace = self.roll_site_session_id L.info( f"pulling prepared tagged_key: {_tagged_key}, packet:{to_one_line_string(packet)}" ) futures.append( RollSite.receive_exeutor_pool.submit(RollSite._thread_receive, self, packet, namespace, roll_site_header)) return futures
def test_ping(self): remote_stub = self.rs_context_remote.stub get_stub = self.rs_context_get.stub req = proxy_pb2.Packet(header=proxy_pb2.Metadata(operator='ping', dst=proxy_pb2.Topic(partyId=self.rs_context_get.party_id))) L.info("starting ping") resp = remote_stub.unaryCall(req) L.info(f"ping finished {resp.header.operator}") self.assertEqual('pong', resp.header.operator) L.info("starting ping reverse") req = proxy_pb2.Packet(header=proxy_pb2.Metadata(operator='ping', dst=proxy_pb2.Topic(partyId=self.rs_context_remote.party_id))) resp = get_stub.unaryCall(req) L.info(f"ping finished {resp.header.operator}") self.assertEqual('pong', resp.header.operator)
def encode_packet(rs_header_inner, batch_inner): header = proxy_pb2.Metadata( src=proxy_pb2.Topic(partyId=rs_header_inner._src_party_id, role=rs_header_inner._src_role), dst=proxy_pb2.Topic(partyId=rs_header_inner._dst_party_id, role=rs_header_inner._dst_role), seq=rs_header_inner._batch_seq, ext=rs_header_inner.to_proto_string(), version=eggroll_version) if batch_inner: result = proxy_pb2.Packet( header=header, body=proxy_pb2.Data(value=batch_inner)) else: result = proxy_pb2.Packet(header=header) return result
def test_get_route_table(self, point='localhost:9470', party='10001', key="shiqili"): channel = grpc.insecure_channel(point) stub = proxy_pb2_grpc.DataTransferServiceStub(channel) salt = self._get_salt() md5hash = hashlib.md5((salt + key).encode("utf-8")).hexdigest() topic_dst = proxy_pb2.Topic(partyId=party) metadata = proxy_pb2.Metadata(dst=topic_dst, operator="get_route_table") data = proxy_pb2.Data(key=md5hash, value=salt.encode('utf-8')) packet = proxy_pb2.Packet(header=metadata, body=data) ret_packet = stub.unaryCall(packet) print(ret_packet.body.value.decode('utf8')) return ret_packet.body.value.decode('utf8')