Args: request: AllocationRequest containing precomputed buffer sizes. Returns: MemoryHandle with RDMA-registered buffer addresses.
(self, request: AllocationRequest)
| 487 | Args: |
| 488 | request: AllocationRequest containing precomputed buffer sizes. |
| 489 | |
| 490 | Returns: |
| 491 | MemoryHandle with RDMA-registered buffer addresses. |
| 492 | """ |
| 493 | buffer_sizes = request.buffer_sizes |
| 494 | room = request.bootstrap_room |
| 495 | self._rdma_buffers[room] = [] |
| 496 | buffers: List[RemoteBuffer] = [] |
| 497 | |
| 498 | for nbytes in buffer_sizes: |
| 499 | if nbytes <= 0: |
| 500 | continue |
| 501 | buf = torch.empty( |
| 502 | (nbytes,), |
| 503 | dtype=torch.uint8, |
| 504 | # device=torch.device(f"cuda:{self.sender_engine_rank}"), |
| 505 | ) |
| 506 | ptr = buf.data_ptr() |
| 507 | self._rdma_buffers[room].append(buf) |
| 508 | buffers.append(RemoteBuffer(addr=ptr, nbytes=nbytes)) |
| 509 | |
| 510 | return MemoryHandle(buffers=buffers) |
| 511 | |
| 512 | def process(self, config): |
| 513 | """ |
| 514 | Generates encoder outputs from prompt and image input. |
| 515 | """ |
| 516 | seed_all(config["seed"]) |
| 517 | self.logger.info("Starting processing in EncoderService...") |
no test coverage detected