# Copyright 2023 The vLLM team.
# Adapted from
# https://github.com/NVIDIA/Megatron-LM/blob/main/megatron/core/parallel_state.py
# Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved.
"""vLLM distributed state.
It takes over the control of the distributed environment from PyTorch.
The typical workflow is:

- call `init_distributed_environment` to initialize the distributed environment.
- call `initialize_model_parallel` or `ensure_model_parallel_initialized` to
 initialize the model parallel groups.

- any code dealing with the distributed stuff

- call `destroy_model_parallel` to destroy the model parallel groups.
- call `destroy_distributed_environment` to destroy the distributed environment.

If you only need to use the distributed environment without model/pipeline
 parallelism, you can skip the model parallel initialization and destruction
 steps.
"""
from typing import Any, Dict, List, Optional, Tuple, Union

import torch
import torch.distributed
from vllm.distributed.parallel_state import (
GroupCoordinator,
_split_tensor_dict,
TensorMetadata,
)

from vllm_mlu.mlu_hijack_utils import MluHijackObject

from examples.cambricon_custom_func.vllm.mlu_hijack.common import init_logger


logger = init_logger(__name__)


def vllm__distributed__GroupCoordinator__send_tensor_dict(
    self,
    tensor_dict: Dict[str, Union[torch.Tensor, Any]],
    dst: Optional[int] = None,
    all_gather_group: Optional["GroupCoordinator"] = None,
) -> Optional[Dict[str, Union[torch.Tensor, Any]]]:
    """Send the input tensor dictionary.
    NOTE: `dst` is the local rank of the source rank.
    """
    # Bypass the function if we are using only 1 GPU.
    if not torch.distributed.is_initialized() or self.world_size == 1:
        return tensor_dict

    all_gather_size = (1 if all_gather_group is None else
                       all_gather_group.world_size)
    all_gather_rank = (0 if all_gather_group is None else
                       all_gather_group.rank_in_group)

    group = self.device_group
    metadata_group = self.cpu_group

    if dst is None:
        dst = (self.rank_in_group + 1) % self.world_size
    assert dst < self.world_size, f"Invalid dst rank ({dst})"

    """
    =============================
    Modifies by vllm_mlu
    =============================
    @brief: Skip send tensor metadata list.
    """
    assert isinstance(
        tensor_dict,
        dict), f"Expecting a dictionary, got {type(tensor_dict)}"
    _, tensor_list = _split_tensor_dict(tensor_dict)
    """
    =============================
    End of MLU Hijack
    =============================
    """
    for tensor in tensor_list:
        if tensor.numel() == 0:
            # Skip sending empty tensors.
            continue

        # send-allgather: send only a slice, then do allgather.
        if (all_gather_group is not None
                and tensor.numel() % all_gather_size == 0):
            tensor = tensor.reshape(all_gather_size, -1)[all_gather_rank]

        if tensor.is_cpu:
            # use metadata_group for CPU tensors
            torch.distributed.send(tensor,
                                   dst=self.ranks[dst],
                                   group=metadata_group)
        else:
            """
            =============================
            Modifies by vllm_mlu
            =============================
            @brief: Modify send to isend.
            """
            # use group for GPU tensors
            torch.distributed.isend(tensor,
                                   dst=self.ranks[dst],
                                   group=group)
            """
            =============================
            End of MLU Hijack
            =============================
            """

    return None

"""
=============================
Modifies by vllm_mlu
=============================
@brief: Add a parameter `recv_metadata_list`.
"""
def vllm__distributed__GroupCoordinator__recv_tensor_dict(
    self,
    src: Optional[int] = None,
    all_gather_group: Optional["GroupCoordinator"] = None,
    recv_metadata_list: List[Tuple[str, Any]] = [],
) -> Optional[Dict[str, Union[torch.Tensor, Any]]]:
    """
    =============================
    End of MLU Hijack
    =============================
    """
    """Recv the input tensor dictionary.
    NOTE: `src` is the local rank of the source rank.
    """
    # Bypass the function if we are using only 1 GPU.
    if not torch.distributed.is_initialized() or self.world_size == 1:
        return None

    all_gather_size = (1 if all_gather_group is None else
                       all_gather_group.world_size)
    all_gather_rank = (0 if all_gather_group is None else
                       all_gather_group.rank_in_group)

    group = self.device_group
    metadata_group = self.cpu_group

    if src is None:
        src = (self.rank_in_group - 1) % self.world_size
    assert src < self.world_size, f"Invalid src rank ({src})"

    """
    =============================
    Modifies by vllm_mlu
    =============================
    @brief: Skip receiving tensor metadata list.
    """
    """
    =============================
    End of MLU Hijack
    =============================
    """
    tensor_dict: Dict[str, Any] = {}
    for key, value in recv_metadata_list:
        if isinstance(value, TensorMetadata):
            tensor = torch.empty(value.size,
                                 dtype=value.dtype,
                                 device=value.device)
            if tensor.numel() == 0:
                # Skip broadcasting empty tensors.
                tensor_dict[key] = tensor
                continue

            # send-allgather: send only a slice, then do allgather.
            use_all_gather = (all_gather_group is not None
                              and tensor.numel() % all_gather_size == 0)

            if use_all_gather:
                orig_shape = tensor.shape
                tensor = tensor.reshape(all_gather_size,
                                        -1)[all_gather_rank]

            if tensor.is_cpu:
                # use metadata_group for CPU tensors
                torch.distributed.recv(tensor,
                                       src=self.ranks[src],
                                       group=metadata_group)
            else:
                """
                =============================
                Modifies by vllm_mlu
                =============================
                @brief: Modify recv to irecv, and wait to finish.
                """
                # use group for GPU tensors
                req = torch.distributed.irecv(tensor,
                                              src=self.ranks[src],
                                              group=group)
                req.wait()
                """
                =============================
                End of MLU Hijack
                =============================
                """
            if use_all_gather:
                # do the allgather
                tensor = all_gather_group.all_gather(  # type: ignore
                    tensor, dim=0)
                tensor = tensor.reshape(orig_shape)

            tensor_dict[key] = tensor
        else:
            tensor_dict[key] = value
    return tensor_dict

MluHijackObject.apply_hijack(
    GroupCoordinator,
    GroupCoordinator.send_tensor_dict,
    vllm__distributed__GroupCoordinator__send_tensor_dict,
)
MluHijackObject.apply_hijack(
    GroupCoordinator,
    GroupCoordinator.recv_tensor_dict,
    vllm__distributed__GroupCoordinator__recv_tensor_dict,
)
