| 48 | |
| 49 | |
| 50 | def update_weights_from_ipc( |
| 51 | zmq_ctx: zmq.Context, |
| 52 | zmq_handle: str, |
| 53 | device_id: int, |
| 54 | *, |
| 55 | run: Callable[[list[tuple[str, torch.Tensor]]], None], |
| 56 | post_hook: Callable[[], None] | None = None, |
| 57 | ): |
| 58 | socket = zmq_ctx.socket(zmq.REP) |
| 59 | socket.connect(zmq_handle) |
| 60 | buffer: torch.Tensor | None = None |
| 61 | device_manager = DeviceManager() |
| 62 | try: |
| 63 | ipc_handle: tuple[Callable, tuple] = socket.recv_pyobj() |
| 64 | assert isinstance(ipc_handle, tuple) |
| 65 | buffer = _rebuild_ipc(ipc_handle, device_id) |
| 66 | assert buffer.dtype == torch.uint8 |
| 67 | socket.send(b"") |
| 68 | except Exception as e: |
| 69 | msg = "".join(traceback.format_exception(type(e), e, e.__traceback__)) |
| 70 | socket.send_string(msg) |
| 71 | socket.recv() # wait for ack |
| 72 | raise |
| 73 | # State machine: |
| 74 | # + receive tensor_metadata -> update_weights |
| 75 | # + receive Exception -> raise and stop |
| 76 | # + receive None first time -> release resources |
| 77 | # + receive None second time -> call post_hook and stop |
| 78 | try: |
| 79 | released = False |
| 80 | while True: |
| 81 | payload: list[FlattenedTensorMetadata] | Exception | None = socket.recv_pyobj() |
| 82 | if released: |
| 83 | assert payload is None, "Should not receive any payload after released" |
| 84 | if post_hook is not None: |
| 85 | post_hook() |
| 86 | device_manager.device_module.synchronize() |
| 87 | socket.send(b"") |
| 88 | break |
| 89 | if payload is None: # done signal |
| 90 | # TODO: wrap all messages to an object instead of None and Exception |
| 91 | device_manager.device_module.synchronize() |
| 92 | released = True |
| 93 | buffer = None |
| 94 | del ipc_handle |
| 95 | |
| 96 | gc.collect() |
| 97 | device_manager.device_module.ipc_collect() |
| 98 | device_manager.device_module.empty_cache() |
| 99 | device_manager.device_module.synchronize() |
| 100 | socket.send(b"") |
| 101 | continue |
| 102 | if isinstance(payload, list): # still updating weights |
| 103 | try: |
| 104 | run(_extract_weights(payload, buffer)) |
| 105 | device_manager.device_module.synchronize() |
| 106 | socket.send(b"") |
| 107 | except Exception as e: # noqa: BLE001 |