# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""
RLHF with FSDP2 training and vLLM expert-parallel inference using **CUDA IPC**
weight transfer and **packed** tensors.

Layout (4 GPUs, TP=1, DP=4, EP):
  * One Ray placement group per GPU.
  * Each PG holds one FSDP training worker and one vLLM ``LLM`` instance
    (sync API) using fractional GPUs so both fit on the same device.
  * The 4 ``LLM`` instances form a DP group via env-var-based SPMD
    coordination (``VLLM_DP_RANK``, ``VLLM_DP_SIZE``, etc.), the same
    mechanism used by ``examples/offline_inference/data_parallel.py``.
  * A ``DataParallelInferenceEngine`` actor spawns all 4 LLM actors,
    waits for initialization, and orchestrates generation / weight-sync.

Uses the built-in ``ray`` send_mode: each FSDP worker calls
``trainer_send_weights`` targeting its colocated LLM actor.

This example was run on 4xH100.
"""

from __future__ import annotations

import os
from dataclasses import asdict

import ray
import torch
import torch.distributed as dist
from huggingface_hub import snapshot_download
from ray.util.placement_group import placement_group
from ray.util.scheduling_strategies import PlacementGroupSchedulingStrategy
from torch.distributed._tensor import DTensor
from torch.distributed.fsdp import fully_shard
from transformers import AutoModelForCausalLM

from vllm import LLM, SamplingParams
from vllm.config import WeightTransferConfig
from vllm.distributed.weight_transfer.ipc_engine import (
    IPCTrainerSendWeightsArgs,
    IPCWeightTransferEngine,
    IPCWeightTransferInitInfo,
)
from vllm.utils.network_utils import get_ip, get_open_port

TRAIN_GPU_FRACTION = float(os.environ.get("RLHF_IPC_TRAIN_GPU_FRACTION", "0.42"))
VLLM_GPU_FRACTION = float(os.environ.get("RLHF_IPC_VLLM_GPU_FRACTION", "0.42"))

MODEL_NAME = "Qwen/Qwen3-30B-A3B"

FSDP_WORLD_SIZE = 4
INFERENCE_TP_SIZE = 1
INFERENCE_DP_SIZE = 4


class MyLLM(LLM):
    """LLM subclass that configures DP env vars for SPMD coordination."""

    def __init__(
        self,
        *args,
        dp_rank: int = 0,
        dp_size: int = 1,
        dp_master_ip: str = "127.0.0.1",
        dp_master_port: int = 0,
        **kwargs,
    ):
        os.environ.pop("CUDA_VISIBLE_DEVICES", None)
        os.environ["VLLM_RAY_PER_WORKER_GPUS"] = str(VLLM_GPU_FRACTION)
        os.environ["VLLM_RAY_BUNDLE_INDICES"] = "0"
        os.environ["VLLM_ALLOW_INSECURE_SERIALIZATION"] = "1"

        os.environ["VLLM_DP_RANK"] = str(dp_rank)
        os.environ["VLLM_DP_RANK_LOCAL"] = str(dp_rank)
        os.environ["VLLM_DP_SIZE"] = str(dp_size)
        os.environ["VLLM_DP_MASTER_IP"] = dp_master_ip
        os.environ["VLLM_DP_MASTER_PORT"] = str(dp_master_port)

        super().__init__(*args, **kwargs)

    def ready(self):
        return True


@ray.remote(num_cpus=0, num_gpus=TRAIN_GPU_FRACTION)
class FSDPTrainWorker:
    """One FSDP2 worker per GPU; colocated with vLLM DP rank via placement group."""

    def __init__(
        self,
        model_name: str,
        rank: int,
        fsdp_world_size: int,
        fsdp_master_addr: str,
        fsdp_master_port: int,
    ):
        self.rank = rank

        os.environ["MASTER_ADDR"] = fsdp_master_addr
        os.environ["MASTER_PORT"] = str(fsdp_master_port)

        dist.init_process_group(backend="nccl", rank=rank, world_size=fsdp_world_size)
        torch.accelerator.set_device_index(0)

        model = AutoModelForCausalLM.from_pretrained(
            model_name, torch_dtype=torch.bfloat16
        )

        self.weight_names = [n for n, _ in model.named_parameters()]
        self.weight_dtype_names = [
            str(p.dtype).split(".")[-1] for _, p in model.named_parameters()
        ]
        self.weight_shapes = [list(p.shape) for _, p in model.named_parameters()]

        for layer in model.model.layers:
            fully_shard(layer)
        fully_shard(model)

        self.model = model

    def get_rank(self):
        return self.rank

    def get_weight_metadata(self):
        return self.weight_names, self.weight_dtype_names, self.weight_shapes

    def gather_and_broadcast_weights_ipc(self, llm_handle, packed: bool = True):
        """All-gather full params; all ranks create IPC handles, rank 0 sends.

        All ranks must call trainer_send_weights so they participate in the
        all_gather_object collective inside _all_gather_and_merge_handles.
        Only rank 0 actually sends the payload to vLLM (gated by _is_rank_zero).
        """

        def _full_param_iter():
            # HF's Qwen3MoeExperts (and other recent HF MoE impls) packs
            # all experts into two fused 3-D tensors per layer:
            #   experts.gate_up_proj  shape (E, 2*I, H)
            #   experts.down_proj     shape (E, H, I)
            # vLLM's Qwen3MoE load_weights still expects the older
            # per-expert HF layout (experts.<i>.gate_proj.weight,
            # experts.<i>.up_proj.weight, experts.<i>.down_proj.weight),
            # so we un-fuse on the fly. Split order matches HF's forward:
            #   gate, up = linear(x, gate_up_proj[i]).chunk(2, dim=-1)
            # → rows [:I] of gate_up_proj[i] are gate, rows [I:] are up.
            params = self.model.state_dict()
            for name in list(params.keys()):
                param = params.pop(name)
                if isinstance(param, DTensor):
                    tensor = param.full_tensor().detach().contiguous()
                else:
                    tensor = param.detach().contiguous()
                del param

                if name.endswith(".experts.gate_up_proj") and tensor.dim() == 3:
                    prefix = name[: -len(".gate_up_proj")]
                    num_experts, two_inter, _ = tensor.shape
                    inter = two_inter // 2
                    for i in range(num_experts):
                        expert = tensor[i]
                        yield (
                            f"{prefix}.{i}.gate_proj.weight",
                            expert[:inter].contiguous(),
                        )
                        yield (
                            f"{prefix}.{i}.up_proj.weight",
                            expert[inter:].contiguous(),
                        )
                    del tensor
                elif name.endswith(".experts.down_proj") and tensor.dim() == 3:
                    prefix = name[: -len(".down_proj")]
                    num_experts = tensor.shape[0]
                    for i in range(num_experts):
                        yield (
                            f"{prefix}.{i}.down_proj.weight",
                            tensor[i].contiguous(),
                        )
                    del tensor
                else:
                    yield name, tensor

        trainer_args = IPCTrainerSendWeightsArgs(
            send_mode="ray",
            llm_handle=llm_handle,
            packed=packed,
            packed_buffer_size_bytes=1024 * 1024 * 1024,  # 1 GB
        )
        IPCWeightTransferEngine.trainer_send_weights(
            iterator=_full_param_iter(),
            trainer_args=trainer_args,
        )


@ray.remote(num_cpus=1)
class DataParallelInferenceEngine:
    """Manages a pool of DP-sharded vLLM LLM actors.

    Spawns one MyLLM actor per placement group, waits for all engines to
    finish initializing, and exposes generation / weight-sync helpers.
    """

    def __init__(
        self,
        model: str,
        pgs: list,
        dp_master_ip: str,
        dp_master_port: int,
    ):
        dp_size = len(pgs)
        self.llm_actors = []
        for r in range(dp_size):
            sched = PlacementGroupSchedulingStrategy(
                placement_group=pgs[r],
                placement_group_capture_child_tasks=True,
            )
            actor = (
                ray.remote(num_cpus=0, num_gpus=0)(MyLLM)
                .options(scheduling_strategy=sched)
                .remote(
                    model=model,
                    enforce_eager=True,
                    tensor_parallel_size=INFERENCE_TP_SIZE,
                    distributed_executor_backend="ray",
                    enable_expert_parallel=True,
                    gpu_memory_utilization=0.35,
                    weight_transfer_config=WeightTransferConfig(backend="ipc"),
                    enable_sleep_mode=True,
                    load_format="dummy",
                    dp_rank=r,
                    dp_size=dp_size,
                    dp_master_ip=dp_master_ip,
                    dp_master_port=dp_master_port,
                )
            )
            self.llm_actors.append(actor)

        ray.get([actor.ready.remote() for actor in self.llm_actors])

    def get_llm_actors(self):
        return self.llm_actors

    def generate(self, prompts: list[str], sampling_params):
        """Distribute prompts round-robin across DP ranks and collect results."""
        dp_size = len(self.llm_actors)
        per_rank: list[list[str]] = [[] for _ in range(dp_size)]
        indices: list[list[int]] = [[] for _ in range(dp_size)]

        for i, prompt in enumerate(prompts):
            rank = i % dp_size
            per_rank[rank].append(prompt)
            indices[rank].append(i)

        refs = [
            actor.generate.remote(per_rank[r], sampling_params)
            for r, actor in enumerate(self.llm_actors)
            if per_rank[r]
        ]
        all_outputs = ray.get(refs)

        ordered = [None] * len(prompts)
        rank_idx = 0
        for r in range(dp_size):
            if per_rank[r]:
                for local_i, orig_i in enumerate(indices[r]):
                    ordered[orig_i] = all_outputs[rank_idx][local_i]
                rank_idx += 1
        return ordered

    def init_weight_transfer(self):
        ray.get(
            [
                actor.init_weight_transfer_engine.remote(
                    dict(init_info=asdict(IPCWeightTransferInitInfo()))
                )
                for actor in self.llm_actors
            ]
        )

    def start_weight_update(self):
        ray.get([actor.start_weight_update.remote() for actor in self.llm_actors])

    def finish_weight_update(self):
        ray.get([actor.finish_weight_update.remote() for actor in self.llm_actors])

    def sleep(self, level: int = 0):
        ray.get([actor.sleep.remote(level=level) for actor in self.llm_actors])

    def wake_up(self, tags: list[str] | None = None):
        ray.get([actor.wake_up.remote(tags=tags) for actor in self.llm_actors])


def main():
    ray.init(
        runtime_env={
            "env_vars": {
                "VLLM_ALLOW_INSECURE_SERIALIZATION": "1",
            }
        }
    )

    assert TRAIN_GPU_FRACTION + VLLM_GPU_FRACTION <= 1.0, (
        "Train + vLLM GPU fractions must sum to at most 1.0 per bundle."
    )

    local_model_path = snapshot_download(MODEL_NAME)
    print(f"[init] Model downloaded to {local_model_path}")

    fsdp_master_addr = get_ip()
    fsdp_master_port = get_open_port()
    dp_master_port = get_open_port()
    dp_master_ip = get_ip()

    # Create one placement group per DP rank (one GPU each).
    pgs = []
    for _ in range(INFERENCE_DP_SIZE):
        pg = placement_group([{"GPU": 1, "CPU": 1}])
        pgs.append(pg)
    ray.get([pg.ready() for pg in pgs])
    print(f"[init] {len(pgs)} placement groups ready.")

    # Launch FSDP training workers, one per PG.
    scheduling = [
        PlacementGroupSchedulingStrategy(
            placement_group=pgs[r],
            placement_group_capture_child_tasks=True,
        )
        for r in range(FSDP_WORLD_SIZE)
    ]

    fsdp_workers = [
        FSDPTrainWorker.options(scheduling_strategy=scheduling[r]).remote(
            local_model_path,
            r,
            FSDP_WORLD_SIZE,
            fsdp_master_addr,
            fsdp_master_port,
        )
        for r in range(FSDP_WORLD_SIZE)
    ]
    ray.get([w.get_rank.remote() for w in fsdp_workers])
    print(f"[init] {FSDP_WORLD_SIZE} FSDP workers ready.")

    # Launch DP inference engine (spawns and initializes all LLM actors).
    inference_engine = DataParallelInferenceEngine.remote(
        model=local_model_path,
        pgs=pgs,
        dp_master_ip=dp_master_ip,
        dp_master_port=dp_master_port,
    )
    llm_actors = ray.get(inference_engine.get_llm_actors.remote())
    print(f"[init] {INFERENCE_DP_SIZE} LLM actors ready.")

    prompts = [
        "Hello, my name is",
        "The president of the United States is",
        "The capital of France is",
        "The future of AI is",
    ]
    sampling_params = SamplingParams(temperature=0)

    print("[generate] Generating with dummy weights...")
    outputs = ray.get(inference_engine.generate.remote(prompts, sampling_params))
    print("-" * 60)
    print("BEFORE weight sync (dummy weights):")
    print("-" * 60)
    for output in outputs:
        print(f"Prompt: {output.prompt!r}")
        print(f"Generated: {output.outputs[0].text!r}")
        print("-" * 60)

    # --- Weight transfer ---
    print("[transfer] Initializing IPC weight transfer...")
    ray.get(inference_engine.init_weight_transfer.remote())

    # Two-phase sleep/wake pattern:
    # 1. sleep(level=1) — offload weights to CPU, discard KV cache
    # 2. wake_up(tags=["weights"]) — bring weights back to GPU (KV cache still free)
    # 3. IPC weight transfer — overwrite weights, plenty of room without KV cache
    # 4. wake_up(tags=["kv_cache"]) — re-allocate KV cache for inference
    print("[sync] Sleeping engines (offload weights + free KV cache)...")
    ray.get(inference_engine.sleep.remote(level=1))

    print("[sync] Waking weights (KV cache stays free)...")
    ray.get(inference_engine.wake_up.remote(tags=["weights"]))

    print("[sync] Starting weight update...")
    ray.get(inference_engine.start_weight_update.remote())

    print("[sync] Packed IPC transfer FSDP → vLLM...")
    ray.get(
        [
            w.gather_and_broadcast_weights_ipc.remote(llm_actors, packed=True)
            for w in fsdp_workers
        ]
    )

    ray.get(inference_engine.finish_weight_update.remote())
    print("[sync] Weight transfer complete.")

    print("[sync] Waking KV cache + scheduling...")
    ray.get(inference_engine.wake_up.remote(tags=["kv_cache", "scheduling"]))

    print("[generate] Generating with synced weights...")
    outputs_updated = ray.get(
        inference_engine.generate.remote(prompts, sampling_params)
    )
    print("-" * 60)
    print("AFTER weight sync (real weights):")
    print("-" * 60)
    for output in outputs_updated:
        print(f"Prompt: {output.prompt!r}")
        print(f"Generated: {output.outputs[0].text!r}")
        print("-" * 60)


if __name__ == "__main__":
    main()