# Copyright 2025 Google LLC
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import asyncio
import copy
import os
from collections import defaultdict
from concurrent.futures import Future
from typing import Dict, List, Optional, Union

import ray
import vllm.envs as envs
from ray.util.placement_group import PlacementGroup
from ray.util.scheduling_strategies import PlacementGroupSchedulingStrategy
from vllm.platforms import current_platform
from vllm.ray.ray_env import get_env_vars_to_copy
from vllm.sequence import IntermediateTensors
from vllm.utils.network_utils import (get_distributed_init_method, get_ip,
                                      get_open_port)
from vllm.v1.core.sched.output import GrammarOutput, SchedulerOutput
from vllm.v1.executor.ray_distributed_executor import \
    RayDistributedExecutor as RayDistributedExecutorV1
from vllm.v1.executor.ray_executor import RayWorkerMetaData
from vllm.v1.executor.ray_utils import RayWorkerWrapper as RayWorkerWrapperV1
from vllm.v1.executor.ray_utils import _wait_until_pg_ready
from vllm.v1.outputs import ModelRunnerOutput

from tpu_inference.distributed.jax_parallel_state import get_pp_group
from tpu_inference.distributed.utils import set_node_kv_ip_port
from tpu_inference.logger import init_logger
from tpu_inference.models.jax.jax_intermediate_tensor import \
    JaxIntermediateTensors
from tpu_inference.runner.tpu_runner import AsyncTPUModelRunnerOutput

logger = init_logger(__name__)


class AsyncResultFuture(Future):

    def __init__(self, result_ids_ref, workers):
        super().__init__()
        self.result_ids_ref = result_ids_ref
        self.workers = workers

    def result(self, timeout=None):
        result_ids = ray.get(self.result_ids_ref, timeout=timeout)
        ret_refs = []
        for result_id, worker in zip(result_ids, self.workers):
            ret_refs.append(
                worker.execute_method.remote("get_execute_model_output",
                                             result_id))
        return ray.get(ret_refs[0], timeout=timeout)


class RayDistributedExecutor(RayDistributedExecutorV1):
    """Ray-based distributed executor for TPU.

    The implementation is similar to vllm/executor/ray_distributed_executor.py
    with these major differences:

    1. self._init_executor():
       VLLM_USE_RAY_SPMD_WORKER=1, in which the driver worker is the same as other workers.
    2. self._initialize_ray_cluster():
       This sets placement_group_specs for TPU.
       In vLLM one GPU maps to one placement group.
       While here one TPU node with all chips maps to one placement group.
    3. self._init_workers_ray():
       This set TPU resources when create each worker.
       And we omit the driver worker related logic.
    """

    @classmethod
    def supports_async_scheduling(cls) -> bool:
        """
        Whether the executor supports async scheduling.
        """
        return True

    def _init_executor(self) -> None:
        self.forward_dag: Optional[ray.dag.CompiledDAG] = None

        os.environ["VLLM_USE_RAY_COMPILED_DAG_CHANNEL_TYPE"] = "shm"

        # Currently, this requires USE_RAY_SPMD_WORKER=True.
        self.use_ray_compiled_dag = True
        # If it is true, then we do not distinguish between the
        # "driver worker" vs other workers. Also, the rank 0 worker will
        # be executed in a remote Ray worker. Currently this requires
        # USE_RAY_COMPILED_DAG=True.
        self.use_ray_spmd_worker = True

        assert self.uses_ray
        self._initialize_ray_cluster()
        placement_group = self.parallel_config.placement_group

        # Disable Ray usage stats collection.
        ray_usage = os.environ.get("RAY_USAGE_STATS_ENABLED", "0")
        if ray_usage != "1":
            os.environ["RAY_USAGE_STATS_ENABLED"] = "0"

        # Create the parallel GPU workers.
        self._init_workers_ray(placement_group)

        self.pp_locks: Optional[List[asyncio.Lock]] = None

        self.scheduler_output: SchedulerOutput | None = None

        # KV connector setup
        self.has_connector = self.vllm_config.kv_transfer_config is not None
        if self.has_connector:
            ip_port = self.collective_rpc("get_node_kv_ip_port")
            for item in ip_port:
                set_node_kv_ip_port(item)
        self.uses_sampler = self.vllm_config.model_config.runner_type != "pooling" and (
            self.vllm_config.ec_transfer_config is None
            or not self.vllm_config.ec_transfer_config.is_ec_producer)

    def _initialize_ray_cluster(self) -> None:
        """Initialize the distributed cluster with Ray.

        it will connect to the Ray cluster and create a placement group
        for the workers, which includes the specification of the resources
        for each distributed worker.

        If a placement group is already provided (e.g., by ray.serve.llm),
        this method will reuse it instead of creating a new one.
        """
        from vllm.platforms import current_platform

        # Check if placement group is already provided (e.g., by ray.serve.llm)
        # This allows ray.serve.llm to pre-create a placement group with the correct
        # TPU topology (e.g., 4 bundles × 4 TPUs for v6e 4x4) and have vLLM reuse it.
        if self.parallel_config.placement_group is not None:
            logger.info(
                f"Using existing placement group: {self.parallel_config.placement_group}"
            )
            return

        if ray.is_initialized():
            logger.info(
                "Ray is already initialized. Skipping Ray initialization.")
        else:
            logger.warning("Ray is not initialized, this is mainly for test.")
            ray.init()

        device_str = current_platform.ray_device_key
        if not device_str:
            raise ValueError(
                f"current platform {current_platform.device_name} does not "
                "support ray.")

        pp_size = self.parallel_config.pipeline_parallel_size
        placement_group_specs: List[Dict[str, float]] = []

        ray_nodes = ray.nodes()
        logger.info(f"RayDistributedExecutor | ray_nodes={ray_nodes}")

        # Filter nodes that have the required TPU
        # This is necessary when the head node doesn't have TPU resources
        # (e.g., KubeRay deployments where head runs on a non-TPU node)
        nodes_with_device = [
            n for n in ray_nodes if device_str in n.get("Resources", {})
        ]
        logger.info(
            f"RayDistributedExecutor | nodes_with_device={len(nodes_with_device)} "
            f"(filtered from {len(ray_nodes)} total nodes)")

        if pp_size == 1:
            placement_group_specs = [{
                device_str: node['Resources'][device_str]
            } for node in nodes_with_device]
        else:
            assert pp_size == len(
                nodes_with_device
            ), f"Cannot use PP across hosts, please set --pipeline-parallel-size to 1 or {len(nodes_with_device)}"
            num_devices_per_pp_rank = self.vllm_config.sharding_config.total_devices
            placement_group_specs = [{
                device_str: num_devices_per_pp_rank
            } for _ in range(pp_size)]

        # vLLM engine is also a worker to execute model with an accelerator,
        # so it requires to have the device in a current node. Check if
        # the current node has at least one device.
        current_ip = get_ip()
        current_node_id = ray.get_runtime_context().get_node_id()
        # Get resources from ray.nodes() which is more reliable
        # than available_resources_per_node() for TPU resource reporting
        current_node_info = next(
            (n for n in ray.nodes() if n["NodeID"] == current_node_id), None)
        current_node_resource = (current_node_info.get("Resources", {})
                                 if current_node_info else {})
        if current_node_resource.get(device_str, 0) < 1:
            raise ValueError(
                f"Current node has no {device_str} available. "
                f"{current_node_resource=}. vLLM engine cannot start without "
                f"{device_str}. Make sure you have at least 1 {device_str} "
                f"available in a node {current_node_id=} {current_ip=}.")
        # This way, at least bundle is required to be created in a current
        # node.
        placement_group_specs[0][f"node:{current_ip}"] = 0.001
        logger.info(
            f"RayDistributedExecutor | placement_group_specs={placement_group_specs}"
        )

        # By default, Ray packs resources as much as possible.
        current_placement_group = ray.util.placement_group(
            placement_group_specs, strategy="PACK")
        _wait_until_pg_ready(current_placement_group)

        assert current_placement_group is not None
        # Set the placement group in the parallel config
        self.parallel_config.placement_group = current_placement_group

    def _init_workers_ray(self, placement_group: "PlacementGroup",
                          **ray_remote_kwargs):
        # The workers are the actual ray actors.
        self.workers: List[RayWorkerWrapper] = []

        # Used in ray compiled DAG: indexed first by PP rank,
        # and then TP rank. In other words, the inner list is
        # the TP group of workers for a PP rank.
        self.pp_tp_workers: List[List[RayWorkerWrapper]] = []

        if self.parallel_config.ray_workers_use_nsight:
            ray_remote_kwargs = self._configure_ray_workers_use_nsight(
                ray_remote_kwargs)

        # Create the workers.
        bundle_indices: List[int]
        if envs.VLLM_RAY_BUNDLE_INDICES:
            # Use the bundle indices specified by the user.
            bundle_indices = list(
                map(int, envs.VLLM_RAY_BUNDLE_INDICES.split(",")))
            assert len(bundle_indices) == self.parallel_config.world_size, \
            ("VLLM_RAY_BUNDLE_INDICES must have the same size"
            f" as the world size, but got {bundle_indices=} "
            f"and {self.parallel_config.world_size=}")
            assert len(set(bundle_indices)) == len(bundle_indices), \
            ("VLLM_RAY_BUNDLE_INDICES cannot have duplicate values,"
            f" but got {bundle_indices=}")
        else:
            bundle_indices = []
            for bundle_id, bundle in enumerate(placement_group.bundle_specs):
                if bundle.get(current_platform.ray_device_key, 0):
                    bundle_indices.append(bundle_id)

        worker_metadata: List[RayWorkerMetaData] = []
        driver_ip = get_ip()
        num_tpu_per_worker = placement_group.bundle_specs[0].get(
            current_platform.ray_device_key, 0)
        for rank, bundle_id in enumerate(bundle_indices):
            scheduling_strategy = PlacementGroupSchedulingStrategy(
                placement_group=placement_group,
                placement_group_capture_child_tasks=True,
                placement_group_bundle_index=bundle_id,
            )
            worker = ray.remote(
                num_cpus=0,
                num_gpus=0,
                resources={
                    current_platform.ray_device_key: num_tpu_per_worker
                },
                scheduling_strategy=scheduling_strategy,
                **ray_remote_kwargs,
            )(RayWorkerWrapper).remote(rpc_rank=rank)
            worker_metadata.append(
                RayWorkerMetaData(worker=worker, created_rank=rank))

        worker_ips = ray.get([
            each.worker.get_node_ip.remote()  # type: ignore[attr-defined]
            for each in worker_metadata
        ])

        for each, ip in zip(worker_metadata, worker_ips):
            each.ip = ip

        logger.debug(f"Initialized worker_metadata: {worker_metadata}")

        ip_counts: Dict[str, int] = {}
        for ip in worker_ips:
            ip_counts[ip] = ip_counts.get(ip, 0) + 1

        def sort_by_driver_then_worker_ip(item: RayWorkerMetaData):
            """
            Sort the workers based on 3 properties:
            1. If the worker is on the same node as the driver (vllm engine),
                it should be placed first.
            2. Then, if the worker is on a node with fewer workers, it should
                be placed first.
            3. Finally, if the work is on a node with smaller IP address, it
                should be placed first.
            """
            ip = item.ip
            return (0 if ip == driver_ip else 1, ip_counts[ip], ip)

        # After sorting, the workers on the same node will be
        # close to each other, and the workers on the driver
        # node will be placed first.
        sorted_worker_metadata = sorted(worker_metadata,
                                        key=sort_by_driver_then_worker_ip)
        start_rank = 0
        for i, item in enumerate(sorted_worker_metadata):
            item.adjusted_rank = i + start_rank
        logger.info(
            f"Initialized sorted worker_metadata: {sorted_worker_metadata}")

        self.workers = [item.worker for item in sorted_worker_metadata]
        rerank_mapping = {
            item.created_rank: item.adjusted_rank
            for item in sorted_worker_metadata
        }
        self.collective_rpc("adjust_rank", args=(rerank_mapping, ))

        # Get the set of TPU IDs used on each node.
        worker_node_and_tpu_ids = []
        for worker in self.workers:
            worker_node_and_tpu_ids.append(
                ray.get(worker.get_node_and_gpu_ids.remote()) \
            ) # type: ignore

        node_workers = defaultdict(list)  # node id -> list of worker ranks
        node_tpus = defaultdict(list)  # node id -> list of tpu ids

        for i, (node_id, tpu_ids) in enumerate(worker_node_and_tpu_ids):
            node_workers[node_id].append(i)
            # `tpu_ids` can be a list of strings or integers.
            # convert them to integers for consistency.
            tpu_ids = [int(x) for x in tpu_ids]
            node_tpus[node_id].extend(tpu_ids)
        for node_id, tpu_ids in node_tpus.items():
            node_tpus[node_id] = sorted(tpu_ids)
        logger.info(
            f"RayDistributedExecutor | node_workers={node_workers} | node_tpus={node_tpus}"
        )

        all_ips = set(worker_ips + [driver_ip])
        n_ips = len(all_ips)
        n_nodes = len(node_workers)

        if n_nodes != n_ips:
            logger.warning(
                f"Got {n_nodes} nodes but with {n_ips} IP addresses. "
                "This is not a typical production setup whose "
                "number of nodes and IPs is euqal. This setup may "
                "lead to unexpected behaviors.")

        # Set environment variables for the driver and workers.
        all_args_to_update_environment_variables = [{
            current_platform.device_control_env_var:
            ",".join(map(str, node_tpus[node_id])),
        } for (node_id, _) in worker_node_and_tpu_ids]

        # Environment variables to copy from driver to workers
        env_vars_to_copy = get_env_vars_to_copy(
            exclude_vars=self.WORKER_SPECIFIC_ENV_VARS,
            additional_vars=set(current_platform.additional_env_vars),
            destination="workers")

        # Copy existing env vars to each worker's args
        for i, args in enumerate(all_args_to_update_environment_variables):
            for name in env_vars_to_copy:
                if name in os.environ:
                    args[name] = os.environ[name]
            logger.debug(
                f"RayDistributedExecutor | Worker {i} environment variables: {args}"
            )

        self._env_vars_for_all_workers = (
            all_args_to_update_environment_variables)

        self.collective_rpc("update_environment_variables",
                            args=(self._get_env_vars_to_be_updated(), ))

        distributed_init_method = get_distributed_init_method(
            driver_ip, get_open_port())

        # Get the Driver's Node ID to identify local vs remote workers
        driver_node_id = ray.get_runtime_context().get_node_id()

        # Initialize the actual workers inside worker wrapper.
        all_kwargs = []
        for rank, (node_id, _) in enumerate(worker_node_and_tpu_ids):
            local_rank = node_workers[node_id].index(rank)
            ip = sorted_worker_metadata[rank].ip
            prev_ip = sorted_worker_metadata[rank - 1].ip if rank > 0 else ""

            worker_vllm_config = self.vllm_config

            # When using object storage (e.g., RunAI), the Leader updates `model` to its local
            # cache path (e.g., /root/.cache/...) during ModelConfig initialization
            # (maybe_pull_model_tokenizer_for_runai), while `model_weights` preserves the original URI after the model is pulled.
            # (Standard HF downloads do not overwrite `model`, allowing workers to pull normally).
            # Since workers on remote nodes cannot access the Leader's filesystem, we create a
            # worker-specific config copy and restore the original GCS URI from `model_weights`.
            # This allows each worker to independently invoke `maybe_pull_model_tokenizer_for_runai`
            # and stream the model from GCS.
            if node_id != driver_node_id and getattr(
                    self.vllm_config, "model_config", None) and getattr(
                        self.vllm_config.model_config, "model_weights", None):
                worker_vllm_config = copy.deepcopy(self.vllm_config)
                worker_vllm_config.model_config.model = worker_vllm_config.model_config.model_weights
                # Unset model_weights so maybe_pull_model_tokenizer_for_runai will pull the model.
                worker_vllm_config.model_config.model_weights = None

            kwargs = dict(
                vllm_config=worker_vllm_config,
                local_rank=local_rank,
                rank=rank,
                distributed_init_method=distributed_init_method,
                is_driver_worker=(not self.parallel_config)
                or (rank % self.parallel_config.tensor_parallel_size == 0),
                ip=ip,
                prev_worker_ip=prev_ip,
            )
            # NOTE(Chenyaaang): Adjust worker's rank to 0 if PP=1.
            # Otherwise if we have 4 ray nodes each with 1 chip and use TP=4,
            # We'll have 4 workers with rank 0,1,2,3 respectively. This
            # contradicts with get_pp_group().
            if self.parallel_config.pipeline_parallel_size == 1:
                kwargs["rank"] = 0
            all_kwargs.append(kwargs)
        self.collective_rpc("init_worker", args=(all_kwargs, ))
        self.collective_rpc("init_device")
        if self.parallel_config.pipeline_parallel_size > 1:
            self.collective_rpc("initialize_pp_transfer_connect")
        self.collective_rpc("load_model")
        if self.use_ray_spmd_worker:
            for pp_rank in range(self.parallel_config.pipeline_parallel_size):
                self.pp_tp_workers.append([])
                num_tp_workers = int(
                    len(self.workers) //
                    self.parallel_config.pipeline_parallel_size)
                for tp_rank in range(num_tp_workers):
                    # PP=2, TP=4, num_tpu_per_worker=2
                    # pp_tp_workers = [[0, 1], [2, 3]]
                    rank = (pp_rank * num_tp_workers) + tp_rank
                    assert len(self.pp_tp_workers[pp_rank]) == tp_rank
                    assert pp_rank < len(self.pp_tp_workers)
                    self.pp_tp_workers[pp_rank].append(self.workers[rank])

    # Ray executor do not need handshake metadata
    # as we pass the kv_parameters through proxy server
    def get_kv_connector_handshake_metadata(self) -> None:
        pass

    # Override the execute_model method to suppprt async scheduling.
    # Once the vLLM V1 Ray executor supports async scheduling natively,
    # we can remove this method.
    def _execute_dag(
        self,
        scheduler_output: SchedulerOutput,
        grammar_output: "GrammarOutput | None",
        non_block: bool = False,
    ) -> ModelRunnerOutput | None | Future[ModelRunnerOutput | None]:
        if not self.scheduler_config.async_scheduling:
            return super()._execute_dag(scheduler_output, grammar_output,
                                        non_block)

        assert non_block
        # Build the compiled DAG for the first time.
        if self.forward_dag is None:  # type: ignore
            self.forward_dag = self._compiled_ray_dag(enable_asyncio=False)

        refs = self.forward_dag.execute(
            (scheduler_output, grammar_output))  # type: ignore
        assert not self.has_connector, "async scheduling with connector not yet supported"
        return AsyncResultFuture(refs, self.workers)


class RayWorkerWrapper(RayWorkerWrapperV1):
    """
    Ray worker wrapper for TPU.

    The implementation is similar to vllm/v1/executor/ray_utils.py
    
    _is_intermediate_tensors: check whether the output is JaxIntermediateTensors.
    _is_last_rank: check whether this Ray worker is the last PP stage.
    """

    def __init__(self, *args, **kwargs) -> None:
        super().__init__(*args, **kwargs)
        #  `_execute_model_outputs` is used to store the actual outputs
        #  for async scheduling, with result_id as the key.
        self._execute_model_outputs = dict()  # type: ignore
        self.result_id = int(0)

    def _is_intermediate_tensors(self, output) -> bool:
        return isinstance(output, JaxIntermediateTensors)

    def _is_last_rank(self) -> bool:
        return get_pp_group().is_last_rank

    # Override the execute_model method to suppprt async scheduling.
    # Once the vLLM V1 Ray executor supports async scheduling natively,
    # we can remove this method.
    def execute_model_ray(
        self,
        execute_model_input: tuple["SchedulerOutput", "GrammarOutput"]
        | tuple["SchedulerOutput", "GrammarOutput", "IntermediateTensors"],
    ) -> Union[
            "ModelRunnerOutput",
            tuple["SchedulerOutput", "GrammarOutput", "IntermediateTensors"],
            int,  # result_id for async scheduling
    ]:
        assert self.vllm_config is not None
        if not self.vllm_config.scheduler_config.async_scheduling:
            return super().execute_model_ray(execute_model_input)

        # This method is used by Ray Compiled Graph to execute the model,
        # and it needs a special logic of self.setup_device_if_necessary()
        self.setup_device_if_necessary()
        assert self.worker is not None, "Worker is not initialized"
        if len(execute_model_input) == 3:
            scheduler_output, grammar_output, intermediate_tensors = (
                execute_model_input)
        else:
            scheduler_output, grammar_output = execute_model_input
            intermediate_tensors = None
        assert self.worker.model_runner is not None
        output = self.worker.model_runner.execute_model(
            scheduler_output, intermediate_tensors)
        assert self._is_last_rank()
        if output is None:
            output = self.worker.model_runner.sample_tokens(grammar_output)
        self.result_id += 1
        self._execute_model_outputs[self.result_id] = output
        return self.result_id

    # Method to get the actual output for async scheduling.
    def get_execute_model_output(self, result_id) -> ModelRunnerOutput:
        assert (self.vllm_config
                and self.vllm_config.scheduler_config.async_scheduling)
        output = self._execute_model_outputs.pop(result_id, None)
        assert output is not None, f"No output found for result_id {result_id}"
        if isinstance(output, AsyncTPUModelRunnerOutput):
            return output.get_output()
        return output
