FATE Eggroll 源码解析

def push(self, obj, parties: list = None, options: dict = None):
    futures = []
    for role_party_id in parties:
        # 本次推送相关的信息保存至 ErRollSiteHeader 中，包括源 party 的 id，目的 party 的 id 等信息

        rs_header = ErRollSiteHeader(
                roll_site_session_id=self.roll_site_session_id, name=self.name, tag=self.tag,
                src_role=self.local_role, src_party_id=self.party_id, dst_role=dst_role, dst_party_id=dst_party_id,
                data_type=data_type)

        # 使用线程池中线程发起实际的数据推送，避免进行阻塞主线程，同时提升并发效率

        # 实际的传输处理是在 RollSiteContext 初始化时注册上的，目前仅注册了 gRPC，因此实际是通过 RollSiteGrpc 发起数据传输

        if isinstance(obj, RollPair):
            future = self._run_thread(self._impl_instance._push_rollpair, obj, rs_header, options)
        else:
            future = self._run_thread(self._impl_instance._push_bytes, obj, rs_header, options)

        futures.append(future)
    return futures

# 对数据进行分片，返回分片序号以及对应的分片数据

def _generate_obj_bytes(py_obj, body_bytes):
    key_id = 0

    # 基于 pickle 进行序列化，保证支持不同类型的数据

    obj_bytes = pickle.dumps(py_obj)
    obj_bytes_len = len(obj_bytes)
    cur_pos = 0


    while cur_pos <= obj_bytes_len:

        # 分片返回数据，每片返回包含 `分片序号, 分片数据`，后续可以基于分片序号进行组装

        yield key_id.to_bytes(int_size, "big"), obj_bytes[cur_pos:cur_pos + body_bytes]
        key_id += 1
        cur_pos += body_bytes

def pair_to_bin_batch(input_iter, limit=None, sendbuf_size=-1):
    pair_count = 0

    # 临时缓冲区

    ba = None
    buffer = None
    writer = None

    # 返回缓冲区内的数据，并清空缓冲区

    def commit(bs=sendbuf_size):
        nonlocal ba
        nonlocal buffer
        nonlocal writer
        bin_batch = None
        if ba:
            bin_batch = bytes(ba[0:writer.get_offset()])

        # 指定缓冲区大小，通过缓存中限制了单次传输的数据量大小
        ba = bytearray(bs)
        buffer = ArrayByteBuffer(ba)
        writer = PairBinWriter(pair_buffer=buffer, data=ba)
        return bin_batch

    # 首次 commit 初始化缓冲区

    commit()

    # 通过持续迭代处理，将原始数据转化为 二进制 batch

    for k, v in input_iter:
        try:
            writer.write(k, v)
            pair_count += 1
            if limit is not None and pair_count >= limit:
                break
        except IndexError as e:

            # 报错时反映当前数据的插入超过缓冲区的大小，通过 commit 将数据返回，并清空缓冲区

            yield commit(max(sendbuf_size, len(k) + len(v) + 1024))
            writer.write(k, v)

    # 将缓冲区内的数据返回

    yield commit()

def _generate_batch_streams(self, pair_iter, batches_per_stream, body_bytes):
        # 将原始 key, value 的迭代器转换为一个自定义格式的自定义字节流格式 batch, 单个 batch 中数据不超过 body_bytes 的大小

        batches = TransferPair.pair_to_bin_batch(pair_iter, sendbuf_size=body_bytes)

        try:
            peek = next(batches)
        except StopIteration as e:
            self._finish_partition = True

        # 将原有的 batch 进行了分组，每组包含 batches_per_stream 个 batch

        def chunk_batch_stream():
            nonlocal self
            nonlocal peek
            cur_batch = peek
            try:
                for i in range(batches_per_stream - 1):
                    next_batch = next(batches)
                    yield cur_batch
                    cur_batch = next_batch
                peek = next(batches)
            except StopIteration as e:
                self._finish_partition = True
            finally:
                yield cur_batch

        while not self._finish_partition:
            self._rs_header._stream_seq += 1

            # 这边使用 yield 返回，因此迭代访问 _generate_batch_streams 获取的每个元素是一个生成器，需要通过 list(xxx) 获取每个元素的列表

            yield chunk_batch_stream()

batch_streams = _generate_batch_streams(pair_iter, batches_per_stream, body_bytes)
for bs in batch_streams:
    for batch in bs:
        print(batch)

def _push_bytes(self, obj, rs_header: ErRollSiteHeader, options: dict = None):

    # 执行数据分片，并将分片的数据转换为 batch 组迭代器

    bs_helper = _BatchStreamHelper(rs_header=rs_header)
    bin_batch_streams = bs_helper._generate_batch_streams(
            pair_iter=_generate_obj_bytes(obj, self.batch_body_bytes),
            batches_per_stream=self.push_batches_per_stream,
            body_bytes=self.batch_body_bytes)

    # 构造 gRPC Client

    grpc_channel_factory = GrpcChannelFactory()
    channel = grpc_channel_factory.create_channel(self.ctx.proxy_endpoint)
    stub = proxy_pb2_grpc.DataTransferServiceStub(channel)

    for batch_stream in bin_batch_streams:
        # 获取单个 batch 组，对应一个生成器，通过 list 方法转换为列表

        batch_stream_data = list(batch_stream)

        # 基于 batch 组构造的 Packet 迭代器，通过 gRPC 进行发送

        stub.push(bs_helper.generate_packet(batch_stream_data, cur_retry), timeout=per_stream_timeout)

def _pull_one(self, rs_header: ErRollSiteHeader, options: dict = None):
    for cur_retry in range(self.pull_max_retry):
        pull_attempts = cur_retry

        # 获取数据传输是否全部完成

        pull_status, all_finished, total_batches, cur_pairs = get_status(self)

        if not all_finished:
            if last_cur_pairs == cur_pairs and cur_pairs > 0:
                raise IOError(f"roll site pull waiting failed because there is no updated progress: rs_key={rs_key}, "
                                f"rs_header={rs_header}, pull_status={pull_status}, last_cur_pairs={last_cur_pairs}, cur_pairs={cur_pairs}")
        else:

            # 数据传输完成，获取数据对应的 RollPair 对象

            rp = self.ctx.rp_ctx.load(name=rp_name, namespace=rp_namespace)

            if data_type == "object":

                # 从 RollPair 中获取完整的对象

                # 基于包序号进行排序，并将排序后的数据组装起来，执行反序列化得到原始数据

                result = _serdes.deserialize(b''.join(map(lambda t: t[1], sorted(rp.get_all(), key=lambda x: int.from_bytes(x[0], "big")))))
            else:
                result = rp

            return result

def gather(self, store):
    client = TransferClient()
    for partition in store._partitions:
        tag = self.__generate_tag(partition._id)
        # 请求地址

        target_endpoint = partition._processor._transfer_endpoint

        # 获取 tag 对应的数据块

        batches = (b.data for b in client.recv(endpoint=target_endpoint, tag=tag, broker=None))

        # 解析数据块，得到 `分片序号, 分片数据` 对
        yield from TransferPair.bin_batch_to_pair(batches)