Function run

tests/test_update.py:135–177

(
    checker_func: callable,
    rank_list: list[list[int]],
    need_error: bool = False,
    expected_exception: Exception | None = None,
    exception_msg: str | None = None,
)

def run(
    checker_func: callable,
    rank_list: list[list[int]],
    need_error: bool = False,
    expected_exception: Exception | None = None,
    exception_msg: str | None = None,
):
    if need_error:
        assert expected_exception is not None, (
            "expected_exception must be provided when need_error is True."
        )
        assert exception_msg is not None, "exception_msg must be provided when need_error is True."
    else:
        assert expected_exception is None, (
            "expected_exception must be None when need_error is False."
        )
        assert exception_msg is None, "exception_msg must be None when need_error is False."

    rank = int(os.getenv("RANK"))
    ctx = get_context("spawn")
    queue = ctx.Queue()
    _device_uuid = _get_physical_gpu_id(device_manager, rank)
    ps = ParameterServer(auto_pg=True)
    _device_uuid = _get_physical_gpu_id(ps.device_manager, rank)
    named_tensors = dict(gen_test_tensors(rank))
    checkpoint_name = "test"
    proc = ctx.Process(target=checker_func, args=(rank, _device_uuid, named_tensors, queue))
    proc.start()
    with pytest.raises(expected_exception) if need_error else nullcontext() as e:
        ps.register_checkpoint(checkpoint_name, named_tensors=named_tensors)
        ps.gather_metas(checkpoint_name)
        for ranks in rank_list:
            ps.update(checkpoint_name, queue.put, ranks=ranks)
            # sleep 3s to wait process group is destroyed
            time.sleep(3)
        if need_error:
            pytest.fail("Test failed: Expected RuntimeError was not raised. Should not reach here.")
    if need_error:
        assert exception_msg in str(e.value)
    ps.unregister_checkpoint(checkpoint_name)
    queue.put(None)
    proc.join()
    assert proc.exitcode == 0


def run_with_files(