Function join

examples/update.py:131–159

(
    ps: ParameterServer,
    checkpoint_name: str,
    load_metas_file: str | None,
    metas_url: str | None,
    req_func: Callable[[list[tuple[str, str]]], None],
    inference_parallel_size: int,
    endpoint: str,
    uds: str | None = None,
)

def join(
    ps: ParameterServer,
    checkpoint_name: str,
    load_metas_file: str | None,
    metas_url: str | None,
    req_func: Callable[[list[tuple[str, str]]], None],
    inference_parallel_size: int,
    endpoint: str,
    uds: str | None = None,
):
    if load_metas_file:
        with open(load_metas_file, "rb") as f:
            metas = _METAS_ADAPTER.validate_json(f.read())
    elif metas_url:
        resp = httpx.get(metas_url, timeout=300.0)
        resp.raise_for_status()
        metas = _METAS_ADAPTER.validate_json(resp.content)
    else:
        raise ValueError("either load_metas_file or metas_url is required")
    ps.init_process_group()
    check_vllm_ready(endpoint, inference_parallel_size, uds)
    dist.barrier()
    with timer("Gather metas before join"):
        ps.gather_metas(checkpoint_name)
    ps.load_metas(metas)
    with timer(
        f"Update weights with setting ranks as range(0, {inference_parallel_size}) by using p2p"
    ):
        ps.update(checkpoint_name, req_func, ranks=list(range(inference_parallel_size)))


if __name__ == "__main__":