Function checker_proc_with_error

tests/test_update.py:52–85

(
    rank: int, device_uuid: str, named_tensors: dict[str, torch.Tensor], queue: Queue
)

def checker_proc_with_error(
    rank: int, device_uuid: str, named_tensors: dict[str, torch.Tensor], queue: Queue
):
    device_manager.device_module.set_device(rank)
    named_tensors = {
        name: tensor.to(device_manager.device_type) for name, tensor in named_tensors.items()
    }
    _ = named_tensors
    _zmq_ctx = zmq.Context()

    def trigger_error(socket_paths: list[tuple[str, str]]):
        socket_paths = dict(socket_paths)
        update_weights_from_ipc(
            _zmq_ctx,
            socket_paths[device_uuid],
            device_id=rank,
            run=error_run,
            post_hook=lambda: device_manager.device_module.synchronize(),
        )

    def error_run(weights: list[tuple[str, torch.Tensor]]):
        _ = weights  # Do some fake processing
        time.sleep(random.uniform(0.1, 0.5))
        if rank == 0:
            raise RuntimeError("Intentional Error for testing.")

    while True:
        socket_paths: list[tuple[str, str]] = queue.get()
        if socket_paths is None:
            break
        try:
            trigger_error(socket_paths)
        except RuntimeError as e:
            assert str(e) == "Some workers failed to update weights"


def checker_proc(rank: int, device_uuid: str, named_tensors: dict[str, torch.Tensor], queue: Queue):