| 46 | |
| 47 | |
| 48 | class GRPCUtils(object): |
| 49 | # default to 256MB |
| 50 | CHUNK_SIZE = ( |
| 51 | int(os.environ["GS_GRPC_CHUNK_SIZE"]) |
| 52 | if "GS_GRPC_CHUNK_SIZE" in os.environ |
| 53 | else 256 * 1024 * 1024 - 1 |
| 54 | ) |
| 55 | |
| 56 | def _generate_chunk_meta(self, chunk): |
| 57 | chunk_meta = attr_value_pb2.ChunkMeta() |
| 58 | chunk_meta.size = len(chunk.buffer) |
| 59 | for k, v in chunk.attr.items(): |
| 60 | chunk_meta.attr[k].CopyFrom(v) |
| 61 | return chunk_meta |
| 62 | |
| 63 | def split(self, dag_def): |
| 64 | """Traverse `large_attr` of op and split into a list of chunks. |
| 65 | |
| 66 | Note that this method will modify `large_attr` attribute of op in dag_def. |
| 67 | |
| 68 | Returns: |
| 69 | Sequence[Sequence[bytes]]: splited chunks. |
| 70 | """ |
| 71 | chunks_list = [] |
| 72 | for op in dag_def.op: |
| 73 | large_attr = attr_value_pb2.LargeAttrValue() |
| 74 | for chunk in op.large_attr.chunk_list.items: |
| 75 | # construct chunk meta |
| 76 | large_attr.chunk_meta_list.items.extend( |
| 77 | [self._generate_chunk_meta(chunk)] |
| 78 | ) |
| 79 | # split buffer |
| 80 | chunks_list.append( |
| 81 | ( |
| 82 | [ |
| 83 | chunk.buffer[i : i + self.CHUNK_SIZE] |
| 84 | for i in range(0, len(chunk.buffer), self.CHUNK_SIZE) |
| 85 | ], |
| 86 | op.key, |
| 87 | ) |
| 88 | ) |
| 89 | # replace chunk with chunk_meta |
| 90 | op.large_attr.CopyFrom(large_attr) |
| 91 | return chunks_list |
| 92 | |
| 93 | def generate_runstep_requests(self, session_id, dag_def): |
| 94 | runstep_requests = [] |
| 95 | chunks_list = self.split(dag_def) |
| 96 | # head |
| 97 | runstep_request = message_pb2.RunStepRequest( |
| 98 | head=message_pb2.RunStepRequestHead(session_id=session_id, dag_def=dag_def) |
| 99 | ) |
| 100 | runstep_requests.append(runstep_request) |
| 101 | # bodies |
| 102 | for chunks, op_key in chunks_list: |
| 103 | for i, chunk in enumerate(chunks): |
| 104 | # check the last element |
| 105 | has_next = True |