Function checker_proc

tests/test_update.py:88–132

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

def checker_proc(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()
    }
    _zmq_ctx = zmq.Context()
    mem_info = device_manager.device_module.mem_get_info()
    memory_usage = mem_info[1] - mem_info[0]
    memory_history: list[int] = [memory_usage]

    def check(names_to_check: dict[str, bool], weights: list[tuple[str, torch.Tensor]]):
        for name, weight in weights:
            if name not in named_tensors:
                continue
            assert (weight == named_tensors[name]).all(), f"Tensor {name} does not match!"
            names_to_check[name] = True

    def check_weights(names_to_check: dict[str, bool], 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=lambda weights: check(names_to_check, weights),
            post_hook=lambda: device_manager.device_module.synchronize(),
        )
        device_manager.device_module.synchronize()
        device_manager.device_module.empty_cache()
        mem_info = device_manager.device_module.mem_get_info()
        memory_usage = mem_info[1] - mem_info[0]
        memory_history.append(memory_usage)
        assert all(names_to_check.values())

    while True:
        socket_paths: list[tuple[str, str]] = queue.get()
        if socket_paths is None:
            break
        names_to_check = dict.fromkeys(named_tensors.keys(), False)
        check_weights(names_to_check, socket_paths)

    mem_info = device_manager.device_module.mem_get_info()
    memory_usage = mem_info[1] - mem_info[0]
    memory_history.append(memory_usage)
    for memory in memory_history[1:]:
        print(f"[rank{rank}] Memory change: {memory - memory_history[0]}")


def run(