ps.py

import argparse
import concurrent.futures
import ctypes
import os
import pickle
import random
import socket
import threading
import time
from collections import defaultdict
from collections.abc import Callable
from datetime import timedelta
from functools import lru_cache
from typing import TYPE_CHECKING, Annotated, Any, BinaryIO, NamedTuple

import httpx
import numpy as np
import torch
import torch.distributed as dist
import zmq
from loguru import logger
from pydantic import BaseModel, PlainSerializer, PlainValidator, WithJsonSchema
from safetensors.torch import safe_open
from torch.multiprocessing.reductions import reduce_tensor


if TYPE_CHECKING:
    from typing_extensions import TypedDict

    class FileMeta(TypedDict):
        key: str  # parameter name
        dtype: torch.dtype
        shape: torch.Size
        type: type
        tp_concat_dim: int


def _dt_validate(value: Any) -> torch.dtype:
    if isinstance(value, str):
        if not value.startswith("torch."):
            raise ValueError(f"dtype {value} should start with torch.")
        try:
            value = getattr(torch, value.split(".")[1])
        except AttributeError as e:
            raise ValueError(f"unknown dtype: {value}") from e
    if not isinstance(value, torch.dtype):
        raise TypeError(f"dtype {value} should be torch.dtype, got {type(value)}")
    return value


_TorchDtype = Annotated[
    torch.dtype,
    PlainValidator(_dt_validate),
    PlainSerializer(lambda x: str(x), return_type=str),
    WithJsonSchema({"type": "string"}, mode="serialization"),
]


def _size_validate(value: Any) -> torch.Size:
    if isinstance(value, list | tuple):
        return torch.Size(value)
    if not isinstance(value, torch.Size):
        raise TypeError(f"size {value} should be torch.Size, got {type(value)}")
    return value


_TorchSize = Annotated[
    torch.Size,
    PlainValidator(_size_validate),
    PlainSerializer(lambda x: tuple(x), return_type=tuple),
    WithJsonSchema({"type": "array", "items": {"type": "integer"}}, mode="serialization"),
]


def _tensor_validate(value: Any) -> torch.Tensor:
    if isinstance(value, torch.Tensor):
        return value
    raise TypeError(f"tensor {value} should be torch.Tensor, got {type(value)}")


_TorchTensor = Annotated[
    torch.Tensor,
    PlainValidator(_tensor_validate),
]


class ParameterMeta(BaseModel):
    name: str
    dtype: _TorchDtype
    shape: _TorchSize


class BucketRange(NamedTuple):
    idx: int  # bucket_idx of MemoryBucket in memory_pool
    offset: int
    size: int


class H2DBucket(BaseModel):
    size: int
    ranges: list[BucketRange]
    items: list[ParameterMeta]


class MemoryBufferMetas(BaseModel):
    metas: list[ParameterMeta]
    ptr: int
    size: int


class MemoryBuffer(BaseModel):
    buffer: _TorchTensor
    size: int
    metas: list[ParameterMeta]


class MemoryBufferMetaList(BaseModel):
    p2p_store_addr: str | None
    memory_buffer_metas_list: list[MemoryBufferMetas]
    rdma_device: str


class DataToGather(MemoryBufferMetaList):
    host_ip: str
    device_uuid: str


# 256 bytes alignment when flatten torch tensors to uint8 buffer
_ALIGN_SIZE = 256


def _align_size(dtype: torch.dtype, shape: torch.Size) -> int:
    return (dtype.itemsize * shape.numel() + _ALIGN_SIZE - 1) // _ALIGN_SIZE * _ALIGN_SIZE


def _to_named_tensor(metas: list[ParameterMeta], offset: int = 0) -> list[dict]:
    ret = []
    for meta in metas:
        size = _align_size(meta.dtype, meta.shape)
        ret.append(
            {
                "name": meta.name,
                "dtype": meta.dtype,
                "shape": meta.shape,
                "offset": offset,
            }
        )
        offset += size
    return ret


def _load_checkpoint_file(file_path: str) -> tuple[int, dict[str, tuple["FileMeta", torch.Tensor]]]:
    def _safetensors_load(fn: str) -> dict[str, tuple["FileMeta", torch.Tensor]]:
        ret = {}
        with safe_open(fn, framework="pt") as f:
            for name in f.keys():  # noqa: SIM118
                weight = f.get_tensor(name)
                meta = {
                    "key": name,
                    "dtype": weight.dtype,
                    "shape": weight.shape,
                    "type": type(weight),
                    "tp_concat_dim": -1,  # safetensors does not support tp_concat_dim
                }
                ret[name] = (meta, weight)
        return ret

    # deprecated, will be removed in the future
    def _fast_np_load(fn: str) -> dict[str, tuple["FileMeta", torch.Tensor]]:
        """load *.np file and return memmap and related tensor meta"""

        def parse_npy_header(fin: BinaryIO) -> dict[str, Any]:
            start = fin.tell()
            major, minor = np.lib.format.read_magic(fin)
            if major == 1 and minor == 0:
                read_header_fn = np.lib.format.read_array_header_1_0
            elif major == 2 and minor == 0:
                read_header_fn = np.lib.format.read_array_header_2_0
            else:
                raise ValueError(
                    f"unknown version {major}.{minor} when parsing npy header from {fn}"
                )
            shape, is_fortran, dtype = read_header_fn(fin)
            return {
                "shape": shape,
                "is_fortran": is_fortran,
                "dtype": dtype,
                "header_length": fin.tell() - start,
            }

        meta_fn = fn + ".meta"
        with open(meta_fn, "rb") as fin:
            meta_lst = pickle.load(fin)

        tensors = []
        offset = 0
        with open(fn, "rb") as fin:
            fin.seek(0, os.SEEK_END)
            filesize = fin.tell()
            fin.seek(0)
            while fin.tell() < filesize:
                tensor_meta = parse_npy_header(fin)
                tensor = np.memmap(
                    fn,
                    dtype=tensor_meta["dtype"],
                    mode="c",
                    offset=offset + tensor_meta["header_length"],
                    shape=tensor_meta["shape"],
                )
                offset += tensor_meta["header_length"] + tensor.nbytes
                fin.seek(offset)
                tensors.append(tensor)

        assert len(meta_lst) == len(tensors)
        ret = {}
        for meta, tensor in zip(meta_lst, tensors):
            if meta["type"] == torch.Tensor:
                tensor = torch.from_numpy(tensor)
            tensor = tensor.view(dtype=meta["dtype"]).view(*meta["shape"])
            ret[meta["key"]] = (meta, tensor)
        return ret

    tp_rank = 0
    if file_path.endswith(".npy"):
        logger.warning("numpy model file is deprecated, will be removed in the future")
        filename_split = os.path.basename(file_path).split(".")
        # if using numpy and want to specify tp rank
        # file should be in model.{layer}.{tp}[.{ep}].npy format
        tp_rank = int(filename_split[2]) if len(filename_split) > 3 else 0
        ret = _fast_np_load(file_path)
    elif file_path.endswith(".safetensors"):
        ret = _safetensors_load(file_path)
    else:
        raise ValueError(f"unsupported file format: {file_path}")
    return tp_rank, ret


def _concat_tp_weights(
    tp_weights: list[torch.Tensor], tp_concat_dim: int, tp_size: int
) -> torch.Tensor:
    """Concat tp weights with meta info.
    If meta.concat_dim is -1, meas this is shared tp weights, just use the first weights.
    Else we will cat weights in concat_dim.
    """
    if tp_concat_dim == -1:
        return tp_weights[0]
    assert tp_size == len(tp_weights)
    if len(tp_weights) == 1:
        return tp_weights[0]
    return torch.cat([w for w in tp_weights], dim=tp_concat_dim)


def _get_physical_gpu_id(device_index: int | None = None) -> str:
    try:
        return f"GPU-{torch.cuda.get_device_properties(device_index).uuid!s}"
    except AssertionError as e:
        raise ValueError(f"fail to get physical gpu id {device_index}") from e


@lru_cache(maxsize=1)
def _get_ip() -> str:
    try:
        # try to get ip from network interface
        with socket.socket(socket.AF_INET, socket.SOCK_DGRAM) as s:
            s.connect(("8.8.8.8", 80))
            return s.getsockname()[0]
    except Exception as e:  # noqa: BLE001
        # fallback to get ip from hostname
        logger.warning(
            f"fail to get ip from network interface, fallback to get ip from hostname: {e}"
        )
        return socket.gethostbyname(socket.gethostname())


def _ibv_get_device_list() -> list[str]:
    lib = ctypes.CDLL("libibverbs.so.1")
    lib.ibv_get_device_list.argtypes = [ctypes.POINTER(ctypes.c_int)]  # int *num_devices
    lib.ibv_get_device_list.restype = ctypes.POINTER(ctypes.c_void_p)  # struct ibv_device **

    lib.ibv_free_device_list.argtypes = [ctypes.POINTER(ctypes.c_void_p)]
    lib.ibv_get_device_name.argtypes = [ctypes.c_void_p]  # struct ibv_device *
    lib.ibv_get_device_name.restype = ctypes.c_char_p  # const char *

    num = ctypes.c_int()
    dev_array = lib.ibv_get_device_list(ctypes.byref(num))
    if not dev_array or num.value <= 0:
        return []

    devices = []
    for i in range(num.value):
        dev_ptr = dev_array[i]  # struct ibv_device *
        name = lib.ibv_get_device_name(dev_ptr)  # const char *
        devices.append(name.decode())
    lib.ibv_free_device_list(dev_array)
    return devices


def _get_rdma_devices() -> list[str]:
    """
    use _ibv_get_device_list to get RDMA devices, if NCCL_IB_HCA has multiple values, just return
    """
    devices_str = os.getenv("PS_P2P_STORE_RDMA_DEVICES")
    if devices_str:
        return devices_str.split(",")
    # if PS_P2P_STORE_RDMA_DEVICES is not set, try to use NCCL_IB_HCA to get RDMA devices
    hca = os.getenv("NCCL_IB_HCA", None)
    return _parse_NCCL_IB_HCA(hca or "", _ibv_get_device_list()) or _ibv_get_device_list()


def _get_my_rdma_device(local_rank: int, gpu_count: int, devices: list[str]) -> str:
    """
    implement network card device allocation, if network card is "mlx5_0,mlx5_1", then 0-3 will share mlx5_0, 4-7 will share mlx5_1, etc.
    """
    if not devices:
        raise RuntimeError("no rdma devices found")
    assert len(devices) <= gpu_count, (
        f"rdma devices count {len(devices)} should be less than or equal to gpu count {gpu_count}"
    )
    assert gpu_count % len(devices) == 0, (
        f"gpu count {gpu_count} should be divisible by rdma devices count {len(devices)}"
    )
    return devices[local_rank // (gpu_count // len(devices))]


def _parse_NCCL_IB_HCA(value: str, available_devices: list[str]) -> list[str]:
    """
    The acceptable value by NCCL_IB_HCA is documented in https://docs.nvidia.com/deeplearning/nccl/user-guide/docs/env.html#id8.
    The Python version parser is referred to the CPP parser in NCCL: https://github.com/NVIDIA/nccl/blob/v2.28.3-1/src/transport/net_ib.cc#L658-L662.

    The list is comma-separated; port numbers are NOT supported yet.
    An optional prefix '^' indicates the list is an exclude list.
    A second optional prefix '=' indicates that the tokens are exact names, otherwise by default NCCL would treat each token as a prefix.
    Please note that when '^' and '=' appear together, only '^=' is allowed, '=^' is not supported.

    Examples:
    - `NCCL_IB_HCA="mlx5"`: Use all cards starting with `mlx5`.
    - `NCCL_IB_HCA="=mlx5_0,mlx5_1"`: Use specific cards `mlx5_0` and `mlx5_1`.
    - `NCCL_IB_HCA="^mlx5"`: Use all cards except those starting with `mlx5`.
    - `NCCL_IB_HCA="^=mlx5_0,mlx5_1"`: Use all cards except `mlx5_0` and `mlx5_1`.
    """
    max_hcas = 32
    if not value or value.strip() == "":
        return available_devices[:max_hcas]

    value = value.strip()
    result = []
    is_exclude = value.startswith("^")
    if is_exclude:
        value = value.removeprefix("^")
    is_exact_match = value.startswith("=")
    if is_exact_match:
        value = value.removeprefix("=")

    device_specs = [spec.strip() for spec in value.split(",") if spec.strip()]

    result = _resolve_device_specs(device_specs, is_exact_match, available_devices)
    if is_exclude:
        result = [dev for dev in available_devices if dev not in result]
    if len(result) > max_hcas:
        result = result[:max_hcas]

    logger.info(f"RDMA Devices from 'NCCL_IB_HCA': {result}")

    return result


def _resolve_device_specs(
    device_specs: list[str], is_exact_match: bool, available_devices: list[str]
) -> list[str]:
    devices = set()
    for spec in device_specs:
        parts = spec.split(":", 1)
        device_name = parts[0].strip()
        # HACK: mooncake transfer engine does not support port specification yet, so we ignore it
        # port = parts[1].strip() if len(parts) > 1 else None
        base_devices = (
            [device_name]
            if device_name in available_devices
            else []
            if is_exact_match
            else [dev for dev in available_devices if dev.startswith(device_name)]
        )

        if not base_devices:
            logger.warning(f"No RDMA device match {device_name=} where {is_exact_match=}.")
            continue

        for base_dev in base_devices:
            devices.add(base_dev)

    return sorted(devices)


def _load_checkpoint(files: list[str]) -> dict[str, torch.Tensor]:
    class TPMeta(BaseModel):
        concat_dim: int
        size: int

    parameters: dict[str, torch.Tensor] = {}
    parameter_metas: dict[str, ParameterMeta] = {}
    tp_metas: dict[str, TPMeta] = {}
    parameters_with_tp: dict[str, dict[int, torch.Tensor]] = {}
    for file in files:
        tp_rank, ret = _load_checkpoint_file(file)
        for parameter_name, (meta, weight) in ret.items():
            if parameter_name not in parameters_with_tp:
                parameters_with_tp[parameter_name] = {}
            parameters_with_tp[parameter_name][tp_rank] = weight
            if parameter_name not in tp_metas:
                tp_metas[parameter_name] = TPMeta(
                    concat_dim=meta["tp_concat_dim"],
                    size=1,
                )
            if parameter_name not in parameter_metas:
                assert isinstance(meta["dtype"], torch.dtype), (
                    f"meta {meta} dtype should be torch.dtype"
                )
                assert isinstance(meta["shape"], torch.Size), (
                    f"meta {meta} shape should be torch.Size"
                )
                parameter_metas[parameter_name] = ParameterMeta(
                    name=parameter_name,
                    shape=meta["shape"],
                    dtype=meta["dtype"],
                )
            tp_meta = tp_metas[parameter_name]
            if tp_meta.concat_dim != -1:
                tp_meta.size = max(tp_meta.size, tp_rank + 1)
    for name, tp_meta in tp_metas.items():
        if tp_meta.concat_dim != -1:
            shape = list(parameter_metas[name].shape)
            shape[tp_meta.concat_dim] = shape[tp_meta.concat_dim] * tp_meta.size
            parameter_metas[name] = ParameterMeta(
                name=name, shape=torch.Size(shape), dtype=parameter_metas[name].dtype
            )
        weights_in_cpu = [parameters_with_tp[name][key] for key in sorted(parameters_with_tp[name])]
        # TODO: here concat is serial, which may be slow
        # but since tp storage is not used in the future
        # we ignore this performance issue for now
        parameters[name] = _concat_tp_weights(weights_in_cpu, tp_meta.concat_dim, tp_meta.size)
    for name, parameter in parameters.items():
        assert name in parameter_metas, f"parameter {name} not found in parameter_metas"
        assert parameter_metas[name].shape == parameter.shape, (
            f"parameter {name} shape mismatch, {parameter_metas[name].shape} != {parameter.shape}"
        )
        assert parameter_metas[name].dtype == parameter.dtype, (
            f"parameter {name} dtype mismatch, {parameter_metas[name].dtype} != {parameter.dtype}"
        )
    return parameters


def _register_checkpoint(
    *,
    files: list[str],
    named_tensors: dict[str, torch.Tensor],
    rank: int | None = None,
) -> list[MemoryBuffer]:
    logger.info(
        f"[rank{rank}] start to register checkpoint with {len(files)} files and {len(named_tensors)} named_tensors"
    )
    if not files and not named_tensors:
        return []
    parameters = _load_checkpoint(files)
    if named_tensors:
        parameters.update(named_tensors)
    bucket_size = max(4 << 30, max(_align_size(x.dtype, x.shape) for x in parameters.values()))

    class MemoryBucket(BaseModel):
        size: int
        metas: list[ParameterMeta]

    buckets: list[MemoryBucket] = [MemoryBucket(size=0, metas=[])]
    for name, tensor in sorted(parameters.items()):
        size = _align_size(tensor.dtype, tensor.shape)
        if buckets[-1].size + size > bucket_size:
            assert buckets[-1], f"buckets[{len(buckets) - 1}] should not be empty"
            buckets.append(MemoryBucket(size=0, metas=[]))
        buckets[-1].metas.append(ParameterMeta(name=name, shape=tensor.shape, dtype=tensor.dtype))
        buckets[-1].size += size

    memory_buffers = [
        MemoryBuffer(buffer=torch.empty(0), size=bucket.size, metas=bucket.metas)
        for bucket in buckets
    ]

    def register_pin_memory(idx: int, size: int) -> tuple[int, torch.Tensor]:
        buffer = torch.empty(size, dtype=torch.uint8, pin_memory=True)
        return idx, buffer

    def register_tensor(buffer: torch.Tensor, offset: int, tensor: torch.Tensor):
        buffer[offset : offset + tensor.nbytes] = tensor.view(-1).view(dtype=torch.uint8)

    with concurrent.futures.ThreadPoolExecutor(max_workers=32) as executor:
        futures = [
            executor.submit(register_pin_memory, idx, bucket.size)
            for idx, bucket in enumerate(buckets)
        ]
        new_futures = []
        for future in concurrent.futures.as_completed(futures):
            idx, buffer = future.result()
            assert buffer.numel() == buckets[idx].size, (
                f"buffer numel {buffer.numel()} should be equal to bucket size {buckets[idx].size}"
            )
            memory_buffers[idx].buffer = buffer
            logger.info(
                f"[rank{rank}] register pin_memory for bucket {idx + 1}/{len(buckets)} finished, "
                f"size {buffer.numel() / 1024 / 1024:.2f}MiB, start to copy tensors to buffer"
            )
            offset = 0
            for meta in buckets[idx].metas:
                name = meta.name
                tensor = parameters[name]
                size = _align_size(tensor.dtype, tensor.shape)
                assert size == _align_size(meta.dtype, meta.shape), (
                    f"tensor {name} size {size} should be equal to meta size {_align_size(meta.dtype, meta.shape)}"
                )
                new_futures.append(executor.submit(register_tensor, buffer, offset, tensor))
                offset += size
        for future in concurrent.futures.as_completed(new_futures):
            future.result()
    return memory_buffers


def request_inference_to_update(
    url: str,
    socket_paths: dict[str, str],
    timeout: float = 300.0,
    uds: str | None = None,
):
    """Send an inference update request to inference server via HTTP or Unix socket.

    Args:
        url (str): The HTTP URL or request path (e.g., "http://localhost:19730/inference") to send the request to.
        socket_paths (dict[str, str]): A dictionary containing device uuid and IPC socket paths for updating weights.
        timeout (float, optional): Request timeout in seconds. Defaults to 300.0.
        uds (str, optional): Path to a Unix domain socket. If provided, the request
            will be sent via the Unix socket instead of HTTP. Defaults to None.

    Raises:
        httpx.HTTPStatusError: If the response contains an HTTP error status.
        httpx.RequestError: If there was an issue while making the request.
    """
    resp = httpx.Client(transport=httpx.HTTPTransport(uds=uds)).post(
        url,
        json={
            "method": "update_weights_from_ipc",
            "args": [socket_paths],
            "timeout": timeout,
        },
        timeout=timeout,
    )
    resp.raise_for_status()


def _gen_h2d_buckets(
    global_metas: dict[int, MemoryBufferMetaList],
    bucket_size: int,
    local_topo: dict[str, set[int]],
    remote_topo: dict[str, set[int]],
    ranks: list[int] | None = None,
) -> list[tuple[int, int, H2DBucket]]:
    buckets: list[tuple[int, H2DBucket]] = []

    for owner_rank, items in global_metas.items():
        buckets.append((owner_rank, H2DBucket(size=0, ranges=[], items=[])))
        for idx, metas in enumerate(items.memory_buffer_metas_list):
            start_offset, offset = 0, 0
            for meta in metas.metas:
                s = _align_size(meta.dtype, meta.shape)
                if buckets[-1][1].size + s > bucket_size:
                    if offset - start_offset > 0:
                        buckets[-1][1].ranges.append(
                            BucketRange(idx, start_offset, offset - start_offset)
                        )
                    start_offset = offset
                    buckets.append((owner_rank, H2DBucket(size=0, ranges=[], items=[])))
                offset += s
                buckets[-1][1].size += s
                buckets[-1][1].items.append(meta)
            buckets[-1][1].ranges.append(BucketRange(idx, start_offset, offset - start_offset))
        assert buckets[-1][1].size > 0, (
            f"buckets[-1][1].size {buckets[-1][1].size} should be greater than 0"
        )
    actual_local_topo = (
        {k: v & set(ranks) for k, v in local_topo.items() if v & set(ranks)}
        if ranks
        else local_topo
    )
    # if ranks is empty, assign the owner_rank as receiver_rank, this is used for colocate architecture
    if not ranks:
        return [(owner_rank, owner_rank, bucket) for owner_rank, bucket in buckets]
    else:
        return _assign_receiver_ranks(buckets, actual_local_topo, remote_topo)


if TYPE_CHECKING:
    from typing import TypeVar

    T = TypeVar("T")


def _assign_receiver_ranks(
    buckets: list[tuple[int, "T"]],
    local_topo: dict[str, set[int]],
    remote_topo: dict[str, set[int]],
) -> list[tuple[int, int, "T"]]:
    """
    (owner_rank, bucket) -> (receiver_rank, owner_rank, bucket)

    Assign receiver ranks to buckets. If ranks is empty, assign the owner_rank as receiver_rank.
    GPU-rdma_device topology will be considered to make full use of the bandwidth.
    """
    if not buckets:
        logger.warning("bucket list is empty, no need to assign receiver ranks")
        return []
    rank_to_rdma_device = {
        rank: rdma_device for rdma_device, ranks in remote_topo.items() for rank in ranks
    }

    # group buckets by owner RDMA devices
    buckets_by_rdma_device = defaultdict(list)
    for owner_rank, bucket in buckets:
        owner_rdma_device = rank_to_rdma_device[owner_rank]
        buckets_by_rdma_device[owner_rdma_device].append((owner_rank, bucket))

    buckets_matrix = list(buckets_by_rdma_device.values())
    assert buckets_matrix, "buckets_matrix should not be empty"

    # Select receiver ranks. We use the minimum rank in each local RDMA device group as receiver rank
    num_receivers = min(len(local_topo), len(buckets_by_rdma_device))
    receiver_list = [min(ranks) for ranks in list(local_topo.values())[:num_receivers]]

    flattened_buckets = [
        buckets_matrix[row][col]
        for col in range(
            max(len(matrix_row) for matrix_row in buckets_matrix) if buckets_matrix else 0
        )
        for row in range(len(buckets_matrix))
        if col < len(buckets_matrix[row])
    ]

    buckets_with_receiver = []
    assigned_cnt = 0
    while assigned_cnt < len(flattened_buckets):
        occupied_devices = set()
        for receiver_rank in receiver_list:
            if assigned_cnt >= len(flattened_buckets):
                break
            owner_rank, bucket = flattened_buckets[assigned_cnt]
            rdma_device = rank_to_rdma_device[owner_rank]
            if rdma_device in occupied_devices:
                break
            buckets_with_receiver.append((receiver_rank, owner_rank, bucket))
            occupied_devices.add(rdma_device)
            assigned_cnt += 1

    return buckets_with_receiver


def _get_master_port(master_port: int | None = None) -> int:
    if master_port is None:
        # HACK: use MASTER_PORT + 1 as master_port, avoid conflict with torchrun's rendezvous port
        # TODO: check whether master_port is available or use a more elegant way
        master_port = int(os.getenv("MASTER_PORT")) + 1
    return master_port


def _get_bcast_rank_map(world_size: int, ranks: list[int] | None) -> dict[int, int]:
    """
    map the real ranks (receiver_rank) to the bcast ranks (0 ~ len(ranks) - 1),
    which are generated in self.init_process_group_for_ranks
    """
    bcast_rank_map: dict[int, int] = {}
    if not ranks:
        bcast_rank_map = {r: r for r in range(world_size)}
    else:
        for i, r in enumerate(ranks):
            bcast_rank_map[r] = i
    return bcast_rank_map


class P2PStore:
    def __init__(self):
        from mooncake.engine import TransferEngine

        self.rank = int(os.getenv("RANK"))
        gpu_count = torch.cuda.device_count()
        local_rank = self.rank % gpu_count
        self.device = _get_my_rdma_device(local_rank, gpu_count, _get_rdma_devices())
        self.ip = _get_ip()

        # we will start at most 8 ps processes, so we use 8 retries to avoid port conflicts in extreme cases
        retry_count = 8
        for i in range(retry_count):
            self.engine = TransferEngine()
            ret = self.engine.initialize(self.ip, "P2PHANDSHAKE", "rdma", self.device)
            if ret == 0:
                break
            # sleep 0.5 ~ 2.0s, to avoid port conflicts when two processes retry at the same time
            sleep_ms = random.randint(500, 2000)
            logger.warning(
                f"[rank{self.rank}] fail to initialize transfer engine, ret {ret}, retry {i + 1}/{retry_count} in {sleep_ms}ms"
            )
            time.sleep(sleep_ms / 1000)
        else:
            raise RuntimeError(f"[rank{self.rank}] fail to initialize transfer engine")
        self.port = self.engine.get_rpc_port()
        self.named_tensors: dict[str, torch.Tensor] = {}
        logger.info(
            f"[rank{self.rank}] p2p store initialized, addr is {self.addr}, rdma device is {self.device}"
        )

    @property
    def addr(self) -> str:
        return f"{self.ip}:{self.port}"

    def register_named_tensors(self, named_tensors: dict[str, torch.Tensor]):
        buffer_addresses = [tensor.data_ptr() for tensor in named_tensors.values()]
        capacities = [tensor.nbytes for tensor in named_tensors.values()]
        self.named_tensors.update(named_tensors)
        for i, name in enumerate(named_tensors.keys()):
            logger.info(
                f"[rank{self.rank}] p2p store register tensor {name} with addr {hex(buffer_addresses[i])} and capacity {capacities[i]}"
            )
        assert self.engine.batch_register_memory(buffer_addresses, capacities) == 0

    def unregister_named_tensors(self, names: list[str]) -> int:
        buffer_addresses = [self.named_tensors[name].data_ptr() for name in names]
        assert self.engine.batch_unregister_memory(buffer_addresses) == 0
        num_unregistered = 0
        for i, name in enumerate(names):
            del self.named_tensors[name]
            logger.info(
                f"[rank{self.rank}] p2p store unregister tensor {name} with addr {hex(buffer_addresses[i])}"
            )
            num_unregistered += 1
        return num_unregistered

    def batch_transfer_sync_read(
        self, target_hostname: str, buf_ptrs: list[int], remote_ptrs: list[int], lens: list[int]
    ):
        assert (
            self.engine.batch_transfer_sync_read(target_hostname, buf_ptrs, remote_ptrs, lens) == 0
        )


class ParameterServer:
    def __init__(
        self,
        *,
        rank: int | None = None,
        world_size: int | None = None,
        auto_pg: bool = False,
        gpu_count: int | None = None,
        mem_fraction: float | None = None,
    ):
        """
        Initialize the parameter server. env RANK, WORLD_SIZE and MASTER_ADDR must be set.

        Args:
            auto_pg: Whether to automatically initialize the process group.
                Notice that if auto_pg is True, will destroy the process group after update.
            mem_fraction: The proportion (as a fraction) of the current free CUDA memory for allocation.
        """
        self._rank = rank or int(os.environ.get("RANK", None))
        self._world_size = world_size or int(os.environ.get("WORLD_SIZE", None))
        self._gpu_count = gpu_count or torch.cuda.device_count()
        self._local_rank = self._rank % self._gpu_count
        self._auto_pg = auto_pg
        self._all_hosts = []
        self._global_device_uuids: list[str] = []
        self._local_rdma_devices: dict[str, set[int]] = defaultdict(set)
        self._remote_rdma_devices: dict[str, set[int]] = defaultdict(set)
        self._mem_fraction = mem_fraction or 0.9

        assert self._rank is not None and self._rank >= 0, self._rank
        assert self._world_size and self._world_size > 0, self._world_size
        assert (
            self._gpu_count is not None
            and self._gpu_count > 0
            and self._gpu_count <= torch.cuda.device_count()
        ), self._gpu_count
        assert (
            self._mem_fraction is not None and self._mem_fraction > 0 and self._mem_fraction <= 1
        ), self._mem_fraction

        self._zmq_ctx = zmq.Context()
        self._zmq_addr_counter = 0

        self._memory_pool: dict[str, list[MemoryBuffer]] = {}
        # dict key is owner_rank, value is a bucket metas list in owner_rank
        self._current_global_parameter_metas: dict[int, MemoryBufferMetaList] = {}
        try:
            self._p2p_store = P2PStore()
        except ImportError as e:
            logger.warning(f"[rank{self._rank}] fail to initialize p2p store due to {e}")
            self._p2p_store = None

        device_index = self._local_rank
        torch.cuda.set_device(device_index)
        self._device_uuid = _get_physical_gpu_id(device_index)
        self._rdma_device = None if self._p2p_store is None else self._p2p_store.device

    def _logger_rank0(self, msg: str):
        if self._local_rank == 0:
            logger.info(msg)

    def get_metas(self) -> dict[int, MemoryBufferMetaList]:
        return self._current_global_parameter_metas

    def load_metas(self, metas: dict[int, MemoryBufferMetaList]):
        self._current_global_parameter_metas = metas
        self._remote_rdma_devices = defaultdict(set)
        for i, meta in self._current_global_parameter_metas.items():
            assert meta.rdma_device is not None, "meta.rdma_device should not be None"
            assert meta.p2p_store_addr is not None, "meta.p2p_store_addr should not be None"
            self._remote_rdma_devices[
                meta.rdma_device + "@" + meta.p2p_store_addr.split(":")[0]
            ].add(i)

    def register_checkpoint(
        self,
        checkpoint_name: str,
        *,
        files: list[str] | None = None,
        named_tensors: dict[str, torch.Tensor] | None = None,
    ) -> None:
        """
        Register a checkpoint to the parameter server. Both files and named_tensors will be registered together.

        Args:
            checkpoint_name: The name of the checkpoint.
            files: The safetensors files to register.
            named_tensors: The named tensors to register.
        """
        try:
            assert checkpoint_name not in self._memory_pool, (
                f"checkpoint {checkpoint_name} already registered"
            )
            self._memory_pool[checkpoint_name] = _register_checkpoint(
                files=files or [], named_tensors=named_tensors or {}, rank=self._rank
            )
            if self._p2p_store is not None:
                self._register_parameters_to_p2p_store(checkpoint_name)
        except Exception:
            logger.exception(
                f"[rank{self._rank}] fail to register checkpoint {checkpoint_name} with files {files}"
            )
            if self._p2p_store is not None:
                self._unregister_parameters_from_p2p_store(checkpoint_name)
            self.unregister_checkpoint(checkpoint_name)
            raise

    def unregister_checkpoint(self, checkpoint_name: str):
        """
        Unregister a checkpoint from the parameter server. This function will also unregister the checkpoint
        from p2p store if p2p store is initialized.
        """
        if checkpoint_name not in self._memory_pool:
            return
        if self._p2p_store is not None:
            num_unregistered = self._unregister_parameters_from_p2p_store(checkpoint_name)
            logger.info(
                f"[rank{self._rank}] unregister {num_unregistered} parameters from p2p store for checkpoint {checkpoint_name}"
            )
        del self._memory_pool[checkpoint_name]
        # see https://github.com/pytorch/pytorch/blob/31d5c675394705f8a6bc767f80ae14bf4f01246b/torch/csrc/cuda/Module.cpp#L2018
        # this works by using torch>=2.5.0
        torch._C._host_emptyCache()

    def gather_metas(self, checkpoint_name: str):
        """
        Gather the parameter metas from all ranks. This will gather memory_buffer, and other metadatas.
        This function should be called before update and init a new value to `self._current_global_parameter_metas`,
        which can be exported by using `self.get_metas` function.
        """
        if self._auto_pg and not dist.is_initialized():
            self.init_process_group()
        assert dist.is_initialized(), "process group is not initialized"
        metas_lst: list[DataToGather | None] = [None for _ in range(self._world_size)]  # type: ignore
        metas = DataToGather(
            memory_buffer_metas_list=[
                MemoryBufferMetas(
                    metas=x.metas,
                    ptr=x.buffer.data_ptr(),
                    size=x.size,
                )
                for x in self._memory_pool.get(checkpoint_name, [])
            ],
            p2p_store_addr=None if self._p2p_store is None else self._p2p_store.addr,
            host_ip=_get_ip(),
            device_uuid=self._device_uuid,
            rdma_device=self._rdma_device or "",
        )

        dist.all_gather_object(metas_lst, metas)

        num_parameters = 0
        all_hosts: list[str] = []
        global_device_uuids: list[str] = []
        for i, metas_buckets in enumerate(metas_lst):
            assert metas_buckets is not None, f"metas_buckets {i} should not be None"
            if i % self._gpu_count == 0 and not self._all_hosts:
                all_hosts.append(metas_buckets.host_ip)
            if not self._global_device_uuids:
                global_device_uuids.append(metas_buckets.device_uuid)
            if metas_buckets.memory_buffer_metas_list:
                self._current_global_parameter_metas[i] = MemoryBufferMetaList(
                    memory_buffer_metas_list=metas_buckets.memory_buffer_metas_list,
                    p2p_store_addr=metas_buckets.p2p_store_addr,
                    rdma_device=metas_buckets.rdma_device,
                )
                num_parameters += sum(len(x.metas) for x in metas_buckets.memory_buffer_metas_list)
            self._local_rdma_devices[
                metas_buckets.rdma_device + "@" + metas_buckets.p2p_store_addr.split(":")[0]
                if metas_buckets.p2p_store_addr
                else metas_buckets.host_ip
            ].add(i)
        if not self._all_hosts:
            self._all_hosts = all_hosts
        if not self._global_device_uuids:
            self._global_device_uuids = global_device_uuids
        # Sender node and Receiver node have the same GPU-rdma_device topology is considered as default.
        # Rewrite the sender's topology (_remote_rdma_devices) by calling load_metas.
        self._remote_rdma_devices = self._local_rdma_devices.copy()
        logger.info(
            f"[rank{self._rank}] gather parameter metas finished, num_parameters: {num_parameters}"
        )

    def init_process_group(
        self,
        *,
        master_addr: str | None = None,
        master_port: int | None = None,
        timeout: timedelta = timedelta(minutes=10),
    ):
        """
        Initialize the process group for the ranks. This global group can be easily destroyed by calling dist.destroy_process_group.

        Args:
            master_port: The specified port of the master node. If not set, will use _get_master_port to get the port.
            timeout: The timeout of the process group.
        """
        master_addr = master_addr or os.getenv("MASTER_ADDR")
        assert master_addr, "master_addr is required"
        store = dist.TCPStore(
            master_addr,
            _get_master_port(master_port),
            self._world_size,
            timeout=timeout,
            is_master=self._rank == 0,
        )
        dist.init_process_group(
            backend="nccl",
            world_size=self._world_size,
            rank=self._rank,
            timeout=timeout,
            store=store,
        )
        logger.info(f"[rank{self._rank}] init process group successfully.")

    def update(
        self,
        checkpoint_name: str,
        req_func: Callable[[list[tuple[str, str]]], None],
        *,
        ranks: list[int] | None = None,
    ) -> None:
        """
        Update the checkpoint to inference engine. This function should be called after gather_metas.

        Args:
            checkpoint_name: The name of the checkpoint.
            req_func: The function to request the inference of inference engine.
            ranks: The ranks to update. If not set, will use fully broadcast to update to all ranks,
                which is the fastest way to update weights, especially in colocated architecture.
                If set, will use p2p to update to the ranks, this is flexible to update to a group of ranks,
                which is useful in disaggregated architecture.
        """
        assert req_func is not None, "req_func is required"
        try:
            # if both ranks is None or [], it will use fully broadcast to update to all ranks
            if not ranks:
                if self._auto_pg and not dist.is_initialized():
                    self.init_process_group()
                self._update_per_bucket(checkpoint_name, req_func)
            else:
                if self._auto_pg:
                    if dist.is_initialized():
                        dist.destroy_process_group()
                        # HACK: wait 2s to ensure destroy is finished
                        time.sleep(2)
                    self.init_process_group_for_ranks(ranks)
                if self._rank not in ranks:
                    return
                self._update_per_bucket(checkpoint_name, req_func, ranks)
            if self._auto_pg:
                dist.destroy_process_group()

            torch.cuda.empty_cache()

            logger.info(
                f"[rank{self._rank}] update checkpoint {checkpoint_name} with ranks {ranks} done. "
                f"Current CUDA allocated {torch.cuda.memory_allocated() / 1024 / 1024} MB, "
                f"reserved {torch.cuda.memory_reserved() / 1024 / 1024} MB."
            )
        except Exception as e:
            logger.exception(
                f"[rank{self._rank}] update checkpoint {checkpoint_name} with ranks {ranks} error {e}"
            )
            raise

    def _bind_zmq_socket(self) -> tuple[zmq.Socket, list[tuple[str, str]]]:
        def zmq_handle(device_uuid: str) -> str:
            return f"ipc://@checkpoint-engine-{device_uuid}-{self._zmq_addr_counter}.sock"

        socket_paths = [(uid, zmq_handle(uid)) for uid in self._global_device_uuids]
        socket = self._zmq_ctx.socket(zmq.REQ)
        socket.bind(zmq_handle(self._device_uuid))
        self._zmq_addr_counter += 1
        return socket, socket_paths

    def _detect_bucket_size(self, *, disable_h2d_buffer: bool = False) -> tuple[int, bool]:
        GiB = 1 << 30  # noqa: N806
        # auto detect bucket size
        tensor = torch.tensor(
            [
                # proportion of current cuda free memory bytes
                int(float(torch.cuda.mem_get_info()[0]) * self._mem_fraction),
                # we use negative value to reuse allreduce min operation
                # for getting the max value of zmq_addr_counter in all ranks
                -self._zmq_addr_counter,
            ],
            dtype=torch.int64,
            device="cuda",
        )
        dist.all_reduce(tensor, op=dist.ReduceOp.MIN)
        tensor = tensor.cpu()
        free_bytes, self._zmq_addr_counter = tensor[0].item(), -tensor[1].item()
        max_tensor_bytes = 0
        for items in self._current_global_parameter_metas.values():
            for metas_list in items.memory_buffer_metas_list:
                for meta in metas_list.metas:
                    max_tensor_bytes = max(max_tensor_bytes, _align_size(meta.dtype, meta.shape))
        free_bytes_divided_3 = free_bytes // (3 * _ALIGN_SIZE) * _ALIGN_SIZE
        if max_tensor_bytes <= free_bytes_divided_3 and not disable_h2d_buffer:
            self._logger_rank0(f"[rank{self._rank}] use h2d buffer")
            # using h2d_buffer can make all ranks' h2d parallel execution
            # the cost is that we need to allocate extra h2d_buffer's GPU memory
            free_bytes = free_bytes_divided_3
        else:
            # if the memory is not enough, it will fallback to disable_h2d_buffer mode,
            # at this time, the bandwidth will be limited by the h2d of a single machine,
            # but we can save GPU memory
            self._logger_rank0(
                f"[rank{self._rank}] disable h2d buffer when max_tensor_bytes {max_tensor_bytes} is larger than free_bytes {free_bytes} // 3"
            )
            free_bytes = free_bytes // (2 * _ALIGN_SIZE) * _ALIGN_SIZE
            assert max_tensor_bytes <= free_bytes, (
                f"max_tensor_bytes {max_tensor_bytes} should be less than free_bytes {free_bytes}"
            )
            disable_h2d_buffer = True
        max_bytes = int(os.getenv("PS_MAX_BUCKET_SIZE_GB", 8)) * GiB
        bucket_size = min(max(max_bytes, max_tensor_bytes), free_bytes)
        logger.info(f"[rank{self._rank}] auto detect bucket size {bucket_size / GiB:.2f} GiB")
        return bucket_size, disable_h2d_buffer

    def _copy_to_buffer(
        self,
        checkpoint_name: str,
        bucket: H2DBucket,
        buffer: torch.Tensor,
        owner_rank: int | None = None,
    ):
        offset = 0
        if owner_rank is not None:
            buf_ptrs, remote_ptrs, lens = [], [], []
            ptr_base = buffer.data_ptr()
            target_addr, ptrs = self._get_addr_ptrs(owner_rank)
        for b in bucket.ranges:
            assert offset + b.size <= bucket.size, (
                f"offset {offset} + size {b.size} > bucket_size {bucket.size}"
            )
            if owner_rank is not None:
                buf_ptrs.append(ptr_base + offset)
                remote_ptrs.append(ptrs[b.idx][0] + b.offset)
                lens.append(b.size)
            else:
                pool = self._memory_pool[checkpoint_name][b.idx]
                buffer[offset : offset + b.size].data.copy_(
                    pool.buffer[b.offset : b.offset + b.size],
                    non_blocking=True,
                )
            offset += b.size
        assert offset == bucket.size, f"offset {offset} != bucket_size {bucket.size}"
        if owner_rank is not None:
            self._p2p_store.batch_transfer_sync_read(target_addr, buf_ptrs, remote_ptrs, lens)
        torch.cuda.synchronize()

    def init_process_group_for_ranks(
        self,
        ranks: list[int],
        *,
        master_port: int | None = None,
        timeout: timedelta = timedelta(minutes=10),
    ):
        """
        Initialize the process group for the ranks. This global group can be easily destroyed by calling dist.destroy_process_group.

        Args:
            ranks: The ranks to initialize the process group. ranks should be a subset of all ranks.
            master_port: The specified port of the master node. If not set, will use _get_master_port to get the port.
            timeout: The timeout of the process group.
        """
        assert not dist.is_initialized()
        assert ranks, "ranks should be set"
        if self._rank not in ranks:
            return
        assert self._all_hosts, "all_hosts should be set"
        assert len(self._all_hosts) == self._world_size // self._gpu_count, (
            f"world_size {self._world_size} should be equal to all_hosts {len(self._all_hosts)}"
        )
        rank = ranks.index(self._rank)
        master_addr = self._all_hosts[ranks[0] // self._gpu_count]
        master_port = _get_master_port(master_port)
        logger.info(
            f"[rank{self._rank}] start to init process group as virtual_rank {rank}, "
            f"master_addr {master_addr}, master_port {master_port}, world_size {len(ranks)}, "
        )
        # only initialize process group and store for ranks, other nodes are not initialized
        # and will not participate in this update. Since they have registered memory addresses
        # to p2p_store at the beginning, update ranks can directly get the memory addresses
        # from other nodes and put the weights into the buffer.
        store = dist.TCPStore(
            master_addr, master_port, len(ranks), is_master=rank == 0, timeout=timeout
        )
        dist.init_process_group(
            backend="nccl", world_size=len(ranks), rank=rank, timeout=timeout, store=store
        )

    def _get_addr_ptrs(self, owner_rank: int) -> tuple[str, list[tuple[int, int]]]:
        addr = self._current_global_parameter_metas[owner_rank].p2p_store_addr
        metas_list = self._current_global_parameter_metas[owner_rank].memory_buffer_metas_list
        return addr, [(metas.ptr, metas.size) for metas in metas_list]

    def _register_parameters_to_p2p_store(self, checkpoint_name: str):
        assert self._p2p_store is not None, "p2p store is not initialized"
        pool = self._memory_pool[checkpoint_name]
        if len(pool) == 0:
            return
        named_tensors, tensor_ptrs = {}, []
        for idx, memory_buffer in enumerate(pool):
            named_tensors[f"memory_pool_{checkpoint_name}_{idx}"] = memory_buffer.buffer
            tensor_ptrs.append((memory_buffer.buffer.data_ptr(), memory_buffer.size))
        self._p2p_store.register_named_tensors(named_tensors)

    def _unregister_parameters_from_p2p_store(self, checkpoint_name: str) -> int:
        assert self._p2p_store is not None, "p2p store is not initialized"
        pool = self._memory_pool[checkpoint_name]
        if len(pool) == 0:
            return 0
        return self._p2p_store.unregister_named_tensors(
            [f"memory_pool_{checkpoint_name}_{idx}" for idx, _ in enumerate(pool)]
        )

    def _update_per_bucket(
        self,
        checkpoint_name: str,
        req_func: Callable[[list[tuple[str, str]]], None],
        ranks: list[int] | None = None,
    ):
        assert len(self._current_global_parameter_metas) != 0, "parameter metas is empty"
        assert dist.is_initialized(), "process group is not initialized"
        # if both ranks is None or [], it will use fully broadcast to update to all ranks
        if not ranks:
            logger.info(f"[rank{self._rank}] update checkpoint {checkpoint_name}")
        # if ranks is set, it will use p2p to update to the ranks
        else:
            assert self._p2p_store is not None, "p2p store is not initialized"
            assert ranks, "ranks should be set"

            need_update = self._rank in ranks
            logger.info(
                f"[rank{self._rank}] update checkpoint {checkpoint_name} p2p, {need_update=} with {ranks=}, "
                f"gpu_count {self._gpu_count}, world_size {self._world_size}"
            )

            if not need_update:
                return
            # first execute a barrier to avoid subsequent cuda oom
            dist.barrier()

        bucket_size, disable_h2d_buffer = self._detect_bucket_size()
        buckets = _gen_h2d_buckets(
            self._current_global_parameter_metas,
            bucket_size,
            self._local_rdma_devices,
            self._remote_rdma_devices,
            ranks,
        )

        h2d_buffer: torch.Tensor | None = (
            None
            if disable_h2d_buffer
            else torch.empty(bucket_size, dtype=torch.uint8, device="cuda")
        )
        # p2p store need to register h2d_buffer to let other ranks read
        if ranks:
            h2d_buffer_name = "__h2d_buffer__"
            if h2d_buffer is not None and self._p2p_store is not None:
                self._p2p_store.register_named_tensors({h2d_buffer_name: h2d_buffer})
        receiver_rank_buckets: list[tuple[int, H2DBucket]] = []
        for receiver_rank, owner_rank, bucket in buckets:
            if receiver_rank != self._rank:
                continue
            receiver_rank_buckets.append((owner_rank, bucket))

        buffer = torch.empty(bucket_size * 2, dtype=torch.uint8, device="cuda")
        handle = reduce_tensor(buffer)

        buckets_by_receiver_rank: dict[int, list[H2DBucket]] = defaultdict(list)
        max_len = 0
        for receiver_rank, _, bucket in buckets:
            buckets_by_receiver_rank[receiver_rank].append(bucket)
            if len(buckets_by_receiver_rank[receiver_rank]) > max_len:
                max_len = len(buckets_by_receiver_rank[receiver_rank])

        socket, socket_paths = self._bind_zmq_socket()
        req_thread = threading.Thread(
            target=req_func,
            args=(socket_paths,),
        )
        req_thread.start()
        socket.send_pyobj(handle)

        gidx = 0
        bcast_rank_map = _get_bcast_rank_map(self._world_size, ranks)
        for i in range(max_len):
            if i < len(receiver_rank_buckets) and not disable_h2d_buffer:
                self._copy_to_buffer(
                    checkpoint_name,
                    receiver_rank_buckets[i][1],
                    h2d_buffer,
                    receiver_rank_buckets[i][0] if ranks else None,
                )
            for receiver_rank, _buckets in buckets_by_receiver_rank.items():
                if i >= len(_buckets):
                    continue
                bucket = _buckets[i]
                alloc, reserved = (
                    torch.cuda.memory_allocated() / 1024 / 1024,
                    torch.cuda.memory_reserved() / 1024 / 1024,
                )
                self._logger_rank0(
                    f"[rank{self._rank}] begin to update bucket {gidx + 1}/{len(buckets)} receiver_rank {receiver_rank} in checkpoint {checkpoint_name}, bucket_size: {bucket.size / 1024 / 1024:.2f}MiB, length: {len(bucket.items)}. "
                    f"Current CUDA allocated {alloc:.2f} MB, "
                    f"reserved {reserved:.2f} MB."
                )
                start = gidx % 2 * bucket_size
                buffer_b: torch.Tensor = buffer[start : start + bucket.size]
                if receiver_rank == self._rank:
                    if disable_h2d_buffer:
                        self._copy_to_buffer(checkpoint_name, bucket, buffer_b)
                    else:
                        buffer_b.data.copy_(h2d_buffer[: bucket.size])
                brank = bcast_rank_map[receiver_rank]
                dist.broadcast(buffer_b, src=brank)
                socket.recv()
                dist.barrier()
                socket.send_pyobj(_to_named_tensor(bucket.items, gidx % 2 * bucket_size))
                gidx += 1

        socket.recv()
        socket.send_pyobj(None)
        socket.recv()
        req_thread.join()
        dist.barrier()
        socket.close()
        if ranks and h2d_buffer is not None:
            self._p2p_store.unregister_named_tensors([h2d_buffer_name])

        torch.cuda.empty_cache()


def _init_api(ps: ParameterServer) -> Any:
    import fastapi
    from fastapi import Request
    from fastapi.responses import JSONResponse, Response

    app = fastapi.FastAPI()

    class RegisterRequest(BaseModel):
        files: list[str]

    class UpdateRequest(BaseModel):
        ranks: list[int] = []
        update_url: str | None = None
        inference_group_ranks: list[int] = []
        timeout: float = 300.0

    def wrap_exception(func: Callable[[], None]) -> Response:
        try:
            func()
        except Exception as e:  # noqa: BLE001
            logger.exception(f"wrap exception {func} failed")
            return JSONResponse(content=str(e), status_code=500)
        return Response(status_code=200)

    @app.post("/v1/checkpoints/{checkpoint_name}/files")
    async def register_files(checkpoint_name: str, req: RegisterRequest, raw: Request) -> Response:
        return wrap_exception(lambda: ps.register_checkpoint(checkpoint_name, files=req.files))

    @app.delete("/v1/checkpoints/{checkpoint_name}")
    async def unregister_checkpoint(checkpoint_name: str) -> Response:
        return wrap_exception(lambda: ps.unregister_checkpoint(checkpoint_name))

    @app.get("/v1/healthz")
    async def healthz() -> Response:
        return Response(status_code=200)

    @app.post("/v1/checkpoints/{checkpoint_name}/gather-metas")
    async def gather_metas(checkpoint_name: str) -> Response:
        return wrap_exception(lambda: ps.gather_metas(checkpoint_name))

    @app.post("/v1/checkpoints/{checkpoint_name}/update")
    async def update(checkpoint_name: str, req: UpdateRequest) -> Response:
        def update_func(socket_paths: list[tuple[str, str]]):
            if req.update_url is None:
                return
            if req.inference_group_ranks:
                socket_paths = [socket_paths[i] for i in req.inference_group_ranks]
            request_inference_to_update(req.update_url, dict(socket_paths), timeout=req.timeout)

        return wrap_exception(lambda: ps.update(checkpoint_name, update_func, ranks=req.ranks))

    return app


@logger.catch(reraise=True)
def run_from_cli():
    import uvicorn

    parser = argparse.ArgumentParser(description="Parameter Server")
    parser.add_argument("--uds", type=str)

    args = parser.parse_args()
    logger.info(
        f"Parameter Server {args=}, master addr: {os.getenv('MASTER_ADDR')}, master port {os.getenv('MASTER_PORT')}"
    )

    assert args.uds and len(args.uds) > 0, args.uds
    ps = ParameterServer(auto_pg=True)
    uvicorn.run(_init_api(ps), uds=args.uds, timeout_keep_alive=60)


if __name__ == "__main__":
    run_from_cli()