# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project

import os

import pytest
import torch
from safetensors.torch import load_file
from torch import nn

from vllm.config import ModelConfig, VllmConfig
from vllm.config.lora import LoRAConfig
from vllm.lora.layers import (
    ColumnParallelLinearWithLoRA,
    MergedColumnParallelLinearWithLoRA,
    ReplicatedLinearWithLoRA,
    RowParallelLinearWithLoRA,
)
from vllm.lora.lora_model import LoRAModel
from vllm.lora.lora_weights import LoRALayerWeights, PackedLoRALayerWeights
from vllm.lora.model_manager import (
    DEFAULT_LANGUAGE_WRAPPER_KEY,
    LoRAMapping,
    LoRAModelManager,
    LRUCacheLoRAModelManager,
)
from vllm.lora.peft_helper import PEFTHelper
from vllm.lora.request import LoRARequest
from vllm.lora.worker_manager import LRUCacheWorkerLoRAManager, WorkerLoRAManager
from vllm.model_executor.layers.fused_moe import GateLinear
from vllm.platforms import current_platform

from .utils import create_peft_lora

EMBEDDING_MODULES = {
    "embed_tokens": "input_embeddings",
    "lm_head": "output_embeddings",
}

DEVICE_TYPE = current_platform.device_type
DEVICES = (
    [f"{DEVICE_TYPE}:{i}" for i in range(min(torch.accelerator.device_count(), 2))]
    if (current_platform.is_cuda_alike() or current_platform.is_xpu())
    else ["cpu"]
)

DEFAULT_DTYPE = torch.get_default_dtype()


@pytest.mark.parametrize("device", DEVICES)
def test_from_lora_tensors(qwen3_lora_files, device):
    tensors = load_file(os.path.join(qwen3_lora_files, "adapter_model.safetensors"))

    peft_helper = PEFTHelper.from_local_dir(
        qwen3_lora_files, max_position_embeddings=4096
    )
    lora_model = LoRAModel.from_lora_tensors(
        1,
        tensors,
        peft_helper=peft_helper,
        device=device,
    )
    for module_name, lora in lora_model.loras.items():
        assert lora.module_name == module_name
        assert lora.rank == 8
        assert lora.lora_alpha == 32
        assert lora.lora_a is not None
        assert lora.lora_b is not None
        assert lora.lora_a.device == torch.device(device)
        assert lora.lora_b.device == torch.device(device)
        assert lora.lora_a.shape[0] == lora.lora_b.shape[1], (
            f"{lora.lora_a.shape=}, {lora.lora_b.shape=}"
        )
        assert lora.lora_a.shape[0] == 8


def create_lora(
    lora_id: int, model: nn.Module, sub_modules: list[str], device: torch.device
) -> LoRAModel:
    loras: dict[str, LoRALayerWeights] = {}
    for name in sub_modules:
        w = model.get_submodule(name).weight
        loras[name] = LoRALayerWeights(
            name,
            8,
            16,
            torch.rand([8, w.shape[1]], device=device),
            torch.rand([w.shape[0], 8], device=device),
        )
    return LoRAModel(lora_id, 8, loras)


def create_packed_lora(
    lora_id: int,
    model: nn.Module,
    module_name,
    replaced_module_names,
    device: torch.device,
    empty_replaced_module_name=None,
) -> LoRAModel:
    w = model.get_submodule(module_name).weight
    loras: dict[str, LoRALayerWeights] = {}
    for replaced_module_name in replaced_module_names:
        if replaced_module_name == empty_replaced_module_name:
            continue
        loras[replaced_module_name] = LoRALayerWeights(
            replaced_module_name,
            8,
            16,
            torch.rand([8, w.shape[1]], device=device),
            torch.rand([w.shape[0] // len(replaced_module_names), 8], device=device),
        )
    return LoRAModel(lora_id, 8, loras)


def test_replace_submodules(default_vllm_config, dist_init, dummy_model):
    model = dummy_model
    manager = LoRAModelManager(
        model,
        1,
        1,
        1,
        LoRAConfig(
            max_lora_rank=8, max_cpu_loras=8, max_loras=8, lora_dtype=DEFAULT_DTYPE
        ),
        torch.device(DEVICES[0]),
        default_vllm_config,
    )
    model = manager.model
    assert isinstance(model.get_submodule("dense1"), ColumnParallelLinearWithLoRA)
    assert isinstance(
        model.get_submodule("layer1.dense1"), ColumnParallelLinearWithLoRA
    )
    assert isinstance(model.get_submodule("dense2"), RowParallelLinearWithLoRA)
    assert isinstance(model.get_submodule("layer1.dense2"), RowParallelLinearWithLoRA)


def test_wrap_replicated_linear_subclasses(default_vllm_config, dist_init, dummy_model):
    from vllm.model_executor.layers.linear import ReplicatedLinear

    class CustomReplicatedLinear(ReplicatedLinear):
        pass

    model = dummy_model
    model.add_module("custom_gate", CustomReplicatedLinear(10, 10, bias=False))

    manager = LoRAModelManager(
        model,
        1,
        1,
        1,
        LoRAConfig(
            max_lora_rank=8, max_cpu_loras=8, max_loras=8, lora_dtype=DEFAULT_DTYPE
        ),
        torch.device(DEVICES[0]),
        default_vllm_config,
    )

    assert isinstance(
        manager.model.get_submodule("custom_gate"), ReplicatedLinearWithLoRA
    )


def test_wrap_gate_linear(default_vllm_config, dist_init, dummy_model):
    model = dummy_model
    model.add_module("router_gate", GateLinear(10, 4, bias=False))

    manager = LoRAModelManager(
        model,
        1,
        1,
        1,
        LoRAConfig(
            max_lora_rank=8, max_cpu_loras=8, max_loras=8, lora_dtype=DEFAULT_DTYPE
        ),
        torch.device(DEVICES[0]),
        default_vllm_config,
    )

    assert isinstance(
        manager.model.get_submodule("router_gate"), ReplicatedLinearWithLoRA
    )


def test_dedup_shared_module_across_paths(default_vllm_config, dist_init, dummy_model):
    """A module reachable from two attribute paths (e.g. a MoE gate held
    both directly on the block and inside its inner runner) must produce a
    single LoRA wrapper. Both paths must end up pointing to that same
    wrapper instance, and only the canonical path should live in
    `manager.modules` — otherwise activate_adapter would call `reset_lora`
    on the alias and clobber weights set under the canonical name.
    """
    from vllm.model_executor.layers.linear import ReplicatedLinear

    class AliasContainer(nn.Module):
        def __init__(self, gate: nn.Module):
            super().__init__()
            self.gate = gate  # canonical path: "moe.gate"

            # Inner submodule holding the SAME gate instance under another
            # path. This mirrors how FusedMoE.runner.gate references the
            # block's gate in qwen3_moe.
            class _Runner(nn.Module):
                def __init__(self, g):
                    super().__init__()
                    self.gate = g  # alias path: "moe.runner.gate"

            self.runner = _Runner(gate)

    gate = ReplicatedLinear(10, 4, bias=False)
    model = dummy_model
    model.add_module("moe", AliasContainer(gate))

    assert model.moe.gate is model.moe.runner.gate

    manager = LoRAModelManager(
        model,
        1,
        1,
        1,
        LoRAConfig(
            max_lora_rank=8, max_cpu_loras=8, max_loras=8, lora_dtype=DEFAULT_DTYPE
        ),
        torch.device(DEVICES[0]),
        default_vllm_config,
    )

    canonical = manager.model.get_submodule("moe.gate")
    alias = manager.model.get_submodule("moe.runner.gate")

    # Same wrapper instance on both paths so forward through either side
    # sees the LoRA-augmented module.
    assert isinstance(canonical, ReplicatedLinearWithLoRA)
    assert alias is canonical

    # Only the canonical path is tracked as a LoRA target. Tracking the
    # alias would cause activate_adapter to reset_lora on it after the
    # canonical entry already populated the weights.
    assert "moe.gate" in manager.modules
    assert "moe.runner.gate" not in manager.modules


def test_lm_head_exempt_from_dedup(default_vllm_config, dist_init, dummy_model):
    """The dedup logic must NOT collapse `lm_head` even when it is reachable
    from another attribute path (tied-embedding models do
    `self.lm_head = self.model.embed_tokens`, sharing the same nn.Module
    instance). The lm_head branch additionally rewires `logits_processor`
    into a `LogitsProcessorWithLoRA`, so skipping it would silently break
    LoRA on lm_head.
    """
    from vllm.lora.layers import LogitsProcessorWithLoRA

    # Add a non-lm_head alias to the same module instance as lm_head. The
    # dedup keys on id(module); without the lm_head exemption the alias
    # would consume the wrapped_by_id slot first and lm_head would be
    # silently skipped, so logits_processor would never be wrapped.
    model = dummy_model
    model.add_module("embed_tokens", model.lm_head)
    assert model.embed_tokens is model.lm_head

    manager = LoRAModelManager(
        model,
        1,
        1,
        1,
        LoRAConfig(
            max_lora_rank=8, max_cpu_loras=8, max_loras=8, lora_dtype=DEFAULT_DTYPE
        ),
        torch.device(DEVICES[0]),
        default_vllm_config,
    )

    # lm_head's special handling still ran: logits_processor got wrapped
    # and the lm_head entry is tracked under self.modules.
    assert isinstance(
        manager.model.get_submodule("logits_processor"), LogitsProcessorWithLoRA
    )
    assert "lm_head" in manager.modules


def test_skip_unsupported_matched_modules(default_vllm_config, dist_init, dummy_model):
    class UnsupportedContainer(nn.Module):
        def __init__(self):
            super().__init__()
            # This name matches a supported target suffix ("dense1"),
            # but nn.Linear is not currently a LoRA-wrappable layer type.
            self.dense1 = nn.Linear(10, 10, bias=False)

    model = dummy_model
    model.add_module("unsupported", UnsupportedContainer())

    manager = LoRAModelManager(
        model,
        1,
        1,
        1,
        LoRAConfig(
            max_lora_rank=8, max_cpu_loras=8, max_loras=8, lora_dtype=DEFAULT_DTYPE
        ),
        torch.device(DEVICES[0]),
        default_vllm_config,
    )

    # Should not crash and should keep unsupported matched modules unchanged.
    assert isinstance(manager.model.get_submodule("unsupported.dense1"), nn.Linear)
    assert "unsupported.dense1" not in manager.modules


def test_target_modules_fail_closed_on_unsupported_matched_modules(
    default_vllm_config, dist_init, dummy_model
):
    class UnsupportedContainer(nn.Module):
        def __init__(self):
            super().__init__()
            self.dense1 = nn.Linear(10, 10, bias=False)

    model = dummy_model
    model.add_module("unsupported", UnsupportedContainer())

    with pytest.raises(ValueError, match="unsupported.dense1"):
        LoRAModelManager(
            model,
            1,
            1,
            1,
            LoRAConfig(
                max_lora_rank=8,
                max_cpu_loras=8,
                max_loras=8,
                lora_dtype=DEFAULT_DTYPE,
                target_modules=["dense1"],
            ),
            torch.device(DEVICES[0]),
            default_vllm_config,
        )


def test_get_dummy_lora_warmup_rank_for_fully_sharded_moe():
    manager = LoRAModelManager.__new__(LoRAModelManager)
    manager.lora_config = LoRAConfig(
        max_lora_rank=64,
        max_cpu_loras=1,
        max_loras=1,
        lora_dtype=DEFAULT_DTYPE,
        fully_sharded_loras=True,
    )

    class DummyModule:
        def __init__(self, tp_size: int, fully_sharded: bool):
            self.tp_size = tp_size
            self.fully_sharded = fully_sharded

    manager.modules = {
        "model.layers.0.self_attn.q_proj": DummyModule(
            tp_size=32,
            fully_sharded=True,
        ),
        "model.layers.0.mlp.experts": DummyModule(
            tp_size=32,
            fully_sharded=True,
        ),
    }

    assert manager.get_dummy_lora_warmup_rank(8) == 32


@pytest.mark.parametrize("device", DEVICES)
def test_lora_model_manager(default_vllm_config, dist_init, dummy_model, device):
    model = dummy_model
    model_lora1 = create_lora(
        1, model, ["layer1.dense1", "dense2", "lm_head"], device=device
    )
    model_lora2 = create_lora(2, model, ["dense1", "dense2", "lm_head"], device=device)
    model_lora3 = create_lora(3, model, ["dense1", "dense2", "lm_head"], device=device)
    manager = LoRAModelManager(
        model,
        2,
        2,
        2,
        LoRAConfig(
            max_lora_rank=8, max_cpu_loras=3, max_loras=2, lora_dtype=DEFAULT_DTYPE
        ),
        device=device,
        vllm_config=default_vllm_config,
    )
    assert all(x is None for x in manager.lora_index_to_id)
    assert manager.add_adapter(model_lora1)
    assert manager.activate_adapter(1)
    assert manager.lora_index_to_id[0] == 1
    assert not manager.add_adapter(model_lora1)
    assert not manager.activate_adapter(1)
    assert manager.add_adapter(model_lora2)
    assert manager.activate_adapter(2)
    assert manager.lora_index_to_id[0] == 1
    assert manager.lora_index_to_id[1] == 2
    assert not manager.add_adapter(model_lora2)
    assert not manager.activate_adapter(2)
    assert manager.add_adapter(model_lora3)
    assert manager.lora_index_to_id[0] == 1
    assert manager.lora_index_to_id[1] == 2
    with pytest.raises(ValueError):
        assert manager.activate_adapter(3)
    assert manager.lora_index_to_id[0] == 1
    assert manager.lora_index_to_id[1] == 2
    assert manager.remove_adapter(model_lora2.id)
    assert manager.lora_index_to_id[1] is None
    assert not manager.remove_adapter(model_lora2.id)
    assert manager.remove_adapter(model_lora1.id)
    assert not manager.remove_adapter(model_lora1.id)
    assert manager.add_adapter(model_lora1)
    assert manager.lora_index_to_id[0] is None
    assert manager.lora_index_to_id[1] is None
    assert manager.add_adapter(model_lora2)
    assert manager.activate_adapter(3)
    assert manager.lora_index_to_id[0] == 3
    assert manager.lora_index_to_id[1] is None
    assert manager.activate_adapter(2)
    assert manager.lora_index_to_id[0] == 3
    assert manager.lora_index_to_id[1] == 2
    assert manager.device == device
    assert (
        manager.punica_wrapper_mapping.get(DEFAULT_LANGUAGE_WRAPPER_KEY).device
        == device
    )
    assert hasattr(manager, "supported_lora_modules")
    assert sorted(manager.supported_lora_modules) == [
        "dense1",
        "dense2",
        "lm_head",
        "output",
    ]


@pytest.mark.parametrize("device", DEVICES)
def test_lora_lru_cache_model_manager(
    default_vllm_config, dist_init, dummy_model, device
):
    model = dummy_model
    model_lora1 = create_lora(
        1, model, ["layer1.dense1", "dense2", "lm_head"], device=device
    )
    model_lora2 = create_lora(2, model, ["dense1", "dense2", "lm_head"], device=device)
    model_lora3 = create_lora(3, model, ["dense1", "dense2", "lm_head"], device=device)
    manager = LRUCacheLoRAModelManager(
        model,
        2,
        2,
        2,
        LoRAConfig(
            max_lora_rank=8, max_cpu_loras=3, max_loras=2, lora_dtype=DEFAULT_DTYPE
        ),
        device=device,
        vllm_config=default_vllm_config,
    )
    assert all(x is None for x in manager.lora_index_to_id)
    assert manager.add_adapter(model_lora1)
    assert manager.activate_adapter(1)
    assert manager.lora_index_to_id[0] == 1
    assert not manager.add_adapter(model_lora1)
    assert not manager.activate_adapter(1)
    assert manager.add_adapter(model_lora2)
    assert manager.activate_adapter(2)
    assert manager.lora_index_to_id[0] == 1
    assert manager.lora_index_to_id[1] == 2
    assert not manager.add_adapter(model_lora2)
    assert not manager.activate_adapter(2)
    assert manager.add_adapter(model_lora3)
    assert manager.lora_index_to_id[0] == 1
    assert manager.lora_index_to_id[1] == 2
    assert manager.activate_adapter(3)
    assert manager.lora_index_to_id[0] == 3
    assert manager.lora_index_to_id[1] == 2
    assert manager.remove_adapter(model_lora2.id)
    assert manager.lora_index_to_id[1] is None
    assert not manager.remove_adapter(model_lora2.id)
    assert manager.remove_adapter(model_lora1.id)
    assert not manager.remove_adapter(model_lora1.id)
    assert manager.add_adapter(model_lora1)
    assert manager.activate_adapter(1)
    assert manager.lora_index_to_id[0] == 3
    assert manager.lora_index_to_id[1] == 1
    assert manager.add_adapter(model_lora2)
    assert manager.deactivate_adapter(3)
    assert manager.lora_index_to_id[0] is None
    assert manager.lora_index_to_id[1] == 1
    assert manager.activate_adapter(2)
    assert manager.lora_index_to_id[0] == 2
    assert manager.lora_index_to_id[1] == 1
    assert manager.activate_adapter(3)
    assert manager.lora_index_to_id[0] == 2
    assert manager.lora_index_to_id[1] == 3
    assert manager.pin_adapter(2)
    assert manager.lora_index_to_id[0] == 2
    assert manager.lora_index_to_id[1] == 3
    assert manager.activate_adapter(1)
    assert manager.lora_index_to_id[0] == 2
    assert manager.lora_index_to_id[1] == 1
    assert manager.deactivate_adapter(2)
    assert manager.lora_index_to_id[0] is None
    assert manager.lora_index_to_id[1] == 1
    assert manager.activate_adapter(3)
    assert manager.lora_index_to_id[0] == 3
    assert manager.lora_index_to_id[1] == 1
    assert manager.pin_adapter(3)
    assert manager.pin_adapter(1)
    with pytest.raises(RuntimeError):
        assert manager.pin_adapter(2)
    assert manager.lora_index_to_id[0] == 3
    assert manager.lora_index_to_id[1] == 1
    with pytest.raises(RuntimeError):
        assert manager.activate_adapter(2)

    assert manager.deactivate_adapter(3)
    assert manager.pin_adapter(2)
    assert manager.lora_index_to_id[0] == 2
    assert manager.lora_index_to_id[1] == 1
    assert manager.remove_adapter(3)
    with pytest.raises(ValueError):
        assert manager.pin_adapter(3)
    assert (
        manager.punica_wrapper_mapping.get(DEFAULT_LANGUAGE_WRAPPER_KEY).device
        == device
    )
    assert manager.device == device


@pytest.mark.parametrize("device", DEVICES)
def test_lru_lora_model_manager(default_vllm_config, dist_init, dummy_model, device):
    # This tests just the LRU cache functionality, everything else is
    # tested in test_lora_model_manager
    model = dummy_model
    model_lora1 = create_lora(
        1, model, ["layer1.dense1", "dense2", "lm_head"], device=device
    )
    model_lora2 = create_lora(2, model, ["dense1", "dense2", "lm_head"], device=device)
    model_lora3 = create_lora(3, model, ["dense1", "dense2", "lm_head"], device=device)
    model_lora4 = create_lora(4, model, ["dense1", "dense2", "lm_head"], device=device)
    manager = LRUCacheLoRAModelManager(
        model,
        2,
        2,
        2,
        LoRAConfig(
            max_lora_rank=8, max_cpu_loras=2, max_loras=2, lora_dtype=DEFAULT_DTYPE
        ),
        device=device,
        vllm_config=default_vllm_config,
    )
    assert all(x is None for x in manager.lora_index_to_id)

    # Add up to capacity
    assert manager.add_adapter(model_lora1)
    assert manager.add_adapter(model_lora2)
    assert manager.activate_adapter(1)
    assert manager.activate_adapter(2)

    assert set(manager.list_adapters()) == {1, 2}
    assert manager.lora_index_to_id[0] == 1
    assert manager.lora_index_to_id[1] == 2

    # Add over capacity
    assert manager.add_adapter(model_lora3)
    assert manager.add_adapter(model_lora4)
    assert manager.activate_adapter(3)
    assert manager.activate_adapter(4)

    assert set(manager.list_adapters()) == {3, 4}
    assert manager.lora_index_to_id[0] == 3
    assert manager.lora_index_to_id[1] == 4

    # Add 3 again to move it to the top and then add 2
    # should return false since it's in already
    assert not manager.add_adapter(model_lora3)
    assert not manager.activate_adapter(3)
    assert manager.add_adapter(model_lora2)
    assert manager.activate_adapter(2)

    assert set(manager.list_adapters()) == {3, 2}
    assert manager.lora_index_to_id[0] == 3
    assert manager.lora_index_to_id[1] == 2

    # Remove manually
    assert manager.remove_adapter(3)
    assert not manager.remove_adapter(3)

    assert set(manager.list_adapters()) == {2}
    assert manager.lora_index_to_id[0] is None
    assert manager.lora_index_to_id[1] == 2

    assert manager.add_adapter(model_lora3)
    assert manager.activate_adapter(3)
    assert manager.add_adapter(model_lora4)
    assert manager.activate_adapter(4)

    assert set(manager.list_adapters()) == {3, 4}
    assert manager.lora_index_to_id[0] == 3
    assert manager.lora_index_to_id[1] == 4

    assert manager.remove_oldest_adapter()
    assert set(manager.list_adapters()) == {4}
    assert manager.lora_index_to_id[0] is None
    assert manager.lora_index_to_id[1] == 4

    assert manager.remove_oldest_adapter()
    assert set(manager.list_adapters()) == set()
    assert all(x is None for x in manager.lora_index_to_id)

    assert not manager.remove_oldest_adapter()
    assert set(manager.list_adapters()) == set()
    assert all(x is None for x in manager.lora_index_to_id)

    # pinning
    assert manager.add_adapter(model_lora3)
    assert manager.activate_adapter(3)
    assert manager.add_adapter(model_lora4)
    assert manager.activate_adapter(4)
    assert set(manager.list_adapters()) == {3, 4}
    with pytest.raises(ValueError):
        assert manager.pin_adapter(1)
    assert manager.pin_adapter(3)
    # Remove manually
    assert manager.remove_adapter(3)
    assert not manager.remove_adapter(3)

    assert set(manager.list_adapters()) == {4}
    assert manager.lora_index_to_id[0] is None
    assert manager.lora_index_to_id[1] == 4

    assert manager.add_adapter(model_lora1)
    assert manager.pin_adapter(1)
    assert manager.add_adapter(model_lora2)
    assert manager.activate_adapter(2)

    assert set(manager.list_adapters()) == {1, 2}
    assert manager.lora_index_to_id[0] == 1
    assert manager.lora_index_to_id[1] == 2

    assert manager.remove_oldest_adapter()
    assert set(manager.list_adapters()) == {1}
    assert manager.lora_index_to_id[0] == 1
    assert manager.lora_index_to_id[1] is None

    with pytest.raises(RuntimeError):
        assert manager.remove_oldest_adapter()

    assert set(manager.list_adapters()) == {1}
    assert (
        manager.punica_wrapper_mapping.get(DEFAULT_LANGUAGE_WRAPPER_KEY).device
        == device
    )
    assert manager.device == device


@pytest.mark.parametrize("device", DEVICES)
def test_lru_cache_worker_adapter_manager(dist_init, dummy_model, device, tmp_path):
    lora_config = LoRAConfig(
        max_lora_rank=8, max_cpu_loras=4, max_loras=4, lora_dtype=DEFAULT_DTYPE
    )

    dummy_lora_files = f"{tmp_path}/lora_adapter"
    os.makedirs(dummy_lora_files, exist_ok=True)
    create_peft_lora(
        dummy_model,
        save_dir=dummy_lora_files,
        target_modules=["layer1.dense1", "dense2"],
        lora_dtype=DEFAULT_DTYPE,
    )

    model_config = ModelConfig(max_model_len=16)
    vllm_config = VllmConfig(model_config=model_config, lora_config=lora_config)

    vllm_config.scheduler_config.max_num_seqs = 4
    vllm_config.scheduler_config.max_num_batched_tokens = 2
    worker_adapter_manager = LRUCacheWorkerLoRAManager(
        vllm_config, device, EMBEDDING_MODULES
    )

    worker_adapter_manager.max_num_seqs = 4
    worker_adapter_manager.max_num_batched_tokens = 2

    worker_adapter_manager.create_lora_manager(dummy_model, vllm_config)

    mapping = LoRAMapping([], [])
    worker_adapter_manager.set_active_adapters(
        [LoRARequest("1", 1, dummy_lora_files), LoRARequest("2", 2, dummy_lora_files)],
        mapping,
    )
    assert worker_adapter_manager.list_adapters() == {1, 2}
    assert worker_adapter_manager._adapter_manager.lora_index_to_id[0] == 1
    assert worker_adapter_manager._adapter_manager.lora_index_to_id[1] == 2

    worker_adapter_manager.set_active_adapters(
        [
            LoRARequest("1", 1, dummy_lora_files),
            LoRARequest("3", 3, dummy_lora_files),
            LoRARequest("4", 4, dummy_lora_files),
        ],
        mapping,
    )
    assert worker_adapter_manager.list_adapters() == {1, 2, 3, 4}
    assert worker_adapter_manager._adapter_manager.lora_index_to_id[0] == 1
    assert worker_adapter_manager._adapter_manager.lora_index_to_id[1] == 2
    assert worker_adapter_manager._adapter_manager.lora_index_to_id[2] == 3
    assert worker_adapter_manager._adapter_manager.lora_index_to_id[3] == 4

    worker_adapter_manager.set_active_adapters(
        [
            LoRARequest("1", 1, dummy_lora_files),
            LoRARequest("2", 2, dummy_lora_files),
            LoRARequest("5", 5, dummy_lora_files),
        ],
        mapping,
    )
    assert worker_adapter_manager.list_adapters() == {1, 2, 4, 5}
    assert worker_adapter_manager._adapter_manager.lora_index_to_id[0] == 1
    assert worker_adapter_manager._adapter_manager.lora_index_to_id[1] == 2
    assert worker_adapter_manager._adapter_manager.lora_index_to_id[2] == 5
    assert worker_adapter_manager._adapter_manager.lora_index_to_id[3] == 4

    worker_adapter_manager.set_active_adapters(
        [
            LoRARequest("1", 1, dummy_lora_files),
            LoRARequest("1", 1, dummy_lora_files),
            LoRARequest("1", 1, dummy_lora_files),
        ],
        mapping,
    )
    assert worker_adapter_manager.list_adapters() == {1, 2, 4, 5}
    assert worker_adapter_manager._adapter_manager.lora_index_to_id[0] == 1
    assert worker_adapter_manager._adapter_manager.lora_index_to_id[1] == 2
    assert worker_adapter_manager._adapter_manager.lora_index_to_id[2] == 5
    assert worker_adapter_manager._adapter_manager.lora_index_to_id[3] == 4

    worker_adapter_manager.set_active_adapters(
        [
            LoRARequest("6", 6, dummy_lora_files),
            LoRARequest("7", 7, dummy_lora_files),
            LoRARequest("8", 8, dummy_lora_files),
        ],
        mapping,
    )
    assert worker_adapter_manager.list_adapters() == {1, 6, 7, 8}
    assert worker_adapter_manager._adapter_manager.lora_index_to_id[0] == 1
    assert worker_adapter_manager._adapter_manager.lora_index_to_id[1] == 7
    assert worker_adapter_manager._adapter_manager.lora_index_to_id[2] == 8
    assert worker_adapter_manager._adapter_manager.lora_index_to_id[3] == 6

    # Over capacity
    with pytest.raises(RuntimeError):
        worker_adapter_manager.set_active_adapters(
            [
                LoRARequest("10", 10, dummy_lora_files),
                LoRARequest("11", 11, dummy_lora_files),
                LoRARequest("12", 12, dummy_lora_files),
                LoRARequest("13", 13, dummy_lora_files),
                LoRARequest("14", 14, dummy_lora_files),
            ],
            mapping,
        )

    assert worker_adapter_manager.device == device
    punica_wrapper = worker_adapter_manager._adapter_manager.punica_wrapper_mapping.get(
        DEFAULT_LANGUAGE_WRAPPER_KEY
    )
    assert punica_wrapper.device == device


@pytest.mark.parametrize("device", DEVICES)
def test_worker_adapter_manager(dist_init, dummy_model_gate_up, device, tmp_path):
    # Should remove every LoRA not specified in the request.
    lora_config = LoRAConfig(
        max_lora_rank=8, max_cpu_loras=4, max_loras=4, lora_dtype=DEFAULT_DTYPE
    )

    model_config = ModelConfig(max_model_len=16)
    vllm_config = VllmConfig(model_config=model_config, lora_config=lora_config)

    vllm_config.scheduler_config.max_num_seqs = 4
    vllm_config.scheduler_config.max_num_batched_tokens = 2

    worker_adapter_manager = WorkerLoRAManager(vllm_config, device, EMBEDDING_MODULES)
    worker_adapter_manager.vocab_size = dummy_model_gate_up.unpadded_vocab_size
    worker_adapter_manager.create_lora_manager(dummy_model_gate_up, vllm_config)

    dummy_lora_files = f"{tmp_path}/lora_adapter"
    os.makedirs(dummy_lora_files, exist_ok=True)
    create_peft_lora(
        dummy_model_gate_up,
        save_dir=dummy_lora_files,
        target_modules=["layer1.dense1", "dense2"],
        lora_dtype=DEFAULT_DTYPE,
    )

    mapping = LoRAMapping([], [])
    worker_adapter_manager.set_active_adapters(
        [LoRARequest("1", 1, dummy_lora_files), LoRARequest("2", 2, dummy_lora_files)],
        mapping,
    )
    assert worker_adapter_manager.list_adapters() == {1, 2}
    assert worker_adapter_manager._adapter_manager.lora_index_to_id[0] == 1
    assert worker_adapter_manager._adapter_manager.lora_index_to_id[1] == 2

    worker_adapter_manager.set_active_adapters(
        [
            LoRARequest("1", 1, dummy_lora_files),
            LoRARequest("3", 3, dummy_lora_files),
            LoRARequest("4", 4, dummy_lora_files),
        ],
        mapping,
    )
    assert worker_adapter_manager.list_adapters() == {1, 3, 4}
    assert worker_adapter_manager._adapter_manager.lora_index_to_id[0] == 1
    assert worker_adapter_manager._adapter_manager.lora_index_to_id[1] == 3
    assert worker_adapter_manager._adapter_manager.lora_index_to_id[2] == 4

    worker_adapter_manager.set_active_adapters(
        [
            LoRARequest("1", 1, dummy_lora_files),
            LoRARequest("2", 2, dummy_lora_files),
            LoRARequest("5", 5, dummy_lora_files),
        ],
        mapping,
    )
    assert worker_adapter_manager.list_adapters() == {1, 2, 5}
    assert worker_adapter_manager._adapter_manager.lora_index_to_id[0] == 1
    assert worker_adapter_manager._adapter_manager.lora_index_to_id[1] == 2
    assert worker_adapter_manager._adapter_manager.lora_index_to_id[2] == 5

    worker_adapter_manager.set_active_adapters(
        [
            LoRARequest("1", 1, dummy_lora_files),
            LoRARequest("1", 1, dummy_lora_files),
            LoRARequest("1", 1, dummy_lora_files),
        ],
        mapping,
    )
    assert worker_adapter_manager.list_adapters() == {1}
    assert worker_adapter_manager._adapter_manager.lora_index_to_id[0] == 1
    assert worker_adapter_manager._adapter_manager.lora_index_to_id[1] is None
    assert worker_adapter_manager._adapter_manager.lora_index_to_id[2] is None

    worker_adapter_manager.set_active_adapters(
        [
            LoRARequest("6", 6, dummy_lora_files),
            LoRARequest("7", 7, dummy_lora_files),
            LoRARequest("8", 8, dummy_lora_files),
        ],
        mapping,
    )
    assert worker_adapter_manager.list_adapters() == {6, 7, 8}
    assert worker_adapter_manager._adapter_manager.lora_index_to_id[0] == 8
    assert worker_adapter_manager._adapter_manager.lora_index_to_id[1] == 6
    assert worker_adapter_manager._adapter_manager.lora_index_to_id[2] == 7

    # Over capacity
    with pytest.raises(RuntimeError):
        worker_adapter_manager.set_active_adapters(
            [
                LoRARequest("10", 10, dummy_lora_files),
                LoRARequest("11", 11, dummy_lora_files),
                LoRARequest("12", 12, dummy_lora_files),
                LoRARequest("13", 13, dummy_lora_files),
                LoRARequest("14", 14, dummy_lora_files),
            ],
            mapping,
        )

    assert worker_adapter_manager.device == device
    punica_wrapper = worker_adapter_manager._adapter_manager.punica_wrapper_mapping.get(
        DEFAULT_LANGUAGE_WRAPPER_KEY
    )
    assert punica_wrapper.device == device


@pytest.mark.parametrize("device", DEVICES)
def test_packed_loras(default_vllm_config, dist_init, dummy_model_gate_up, device):
    model = dummy_model_gate_up
    model_lora = create_packed_lora(
        1,
        model,
        module_name="gate_up_proj",
        replaced_module_names=["gate_proj", "up_proj"],
        device=device,
    )
    model_lora1 = create_packed_lora(
        2,
        model,
        module_name="gate_up_proj",
        replaced_module_names=["gate_proj", "up_proj"],
        device=device,
        empty_replaced_module_name="gate_proj",
    )

    manager = LoRAModelManager(
        model,
        2,
        2,
        2,
        LoRAConfig(
            max_lora_rank=8, max_cpu_loras=2, max_loras=2, lora_dtype=DEFAULT_DTYPE
        ),
        device=device,
        vllm_config=default_vllm_config,
    )
    model = manager.model

    assert isinstance(
        model.get_submodule("gate_up_proj"), MergedColumnParallelLinearWithLoRA
    )
    # Verify packed lora is correct
    model_lora_clone = model_lora.clone(1)
    model_lora_clone1 = model_lora1.clone(1)
    assert manager.add_adapter(model_lora)
    assert manager.add_adapter(model_lora1)

    assert model_lora.get_lora("gate_proj") is None
    assert model_lora.get_lora("up_proj") is None
    assert model_lora1.get_lora("up_proj") is None
    packed_lora = model_lora.get_lora("gate_up_proj")
    assert packed_lora and isinstance(packed_lora, PackedLoRALayerWeights)

    torch.testing.assert_close(
        packed_lora.lora_a[0], model_lora_clone.get_lora("gate_proj").lora_a
    )
    torch.testing.assert_close(
        packed_lora.lora_b[0], model_lora_clone.get_lora("gate_proj").lora_b
    )
    torch.testing.assert_close(
        packed_lora.lora_a[1], model_lora_clone.get_lora("up_proj").lora_a
    )
    torch.testing.assert_close(
        packed_lora.lora_b[1], model_lora_clone.get_lora("up_proj").lora_b
    )

    packed_lora1 = model_lora1.get_lora("gate_up_proj")
    assert packed_lora1 and isinstance(packed_lora1, PackedLoRALayerWeights)

    assert packed_lora1.lora_a[0] is None
    assert packed_lora1.lora_b[0] is None
    torch.testing.assert_close(
        packed_lora1.lora_a[1], model_lora_clone1.get_lora("up_proj").lora_a
    )
    torch.testing.assert_close(
        packed_lora1.lora_b[1], model_lora_clone1.get_lora("up_proj").lora_b
    )


def _test_target_modules(
    model,
    target_modules: list[str] | None,
    device: str,
    expected_lora: list[tuple[str, type]],
    expected_no_lora: list[tuple[str, type]],
    vllm_config,
):
    """Create a LoRAModelManager and assert which modules have LoRA applied."""
    LoRAModelManager(
        model,
        2,
        2,
        2,
        LoRAConfig(
            max_lora_rank=8,
            max_cpu_loras=2,
            max_loras=2,
            lora_dtype=DEFAULT_DTYPE,
            target_modules=target_modules,
        ),
        device=device,
        vllm_config=vllm_config,
    )
    for module_path, lora_cls in expected_lora:
        assert isinstance(model.get_submodule(module_path), lora_cls)
    for module_path, lora_cls in expected_no_lora:
        assert not isinstance(model.get_submodule(module_path), lora_cls)


@pytest.mark.parametrize("device", DEVICES)
def test_target_modules_config(default_vllm_config, dist_init, dummy_model, device):
    """Test that target_modules config restricts which modules get LoRA applied."""
    _test_target_modules(
        dummy_model,
        ["dense1"],
        device,
        expected_lora=[
            ("dense1", ColumnParallelLinearWithLoRA),
            ("layer1.dense1", ColumnParallelLinearWithLoRA),
        ],
        expected_no_lora=[
            ("dense2", RowParallelLinearWithLoRA),
            ("layer1.dense2", RowParallelLinearWithLoRA),
        ],
        vllm_config=default_vllm_config,
    )


@pytest.mark.parametrize("device", DEVICES)
def test_target_modules_multiple(default_vllm_config, dist_init, dummy_model, device):
    """Test that multiple target_modules work correctly."""
    _test_target_modules(
        dummy_model,
        ["dense1", "dense2"],
        device,
        expected_lora=[
            ("dense1", ColumnParallelLinearWithLoRA),
            ("layer1.dense1", ColumnParallelLinearWithLoRA),
            ("dense2", RowParallelLinearWithLoRA),
            ("layer1.dense2", RowParallelLinearWithLoRA),
        ],
        expected_no_lora=[],
        vllm_config=default_vllm_config,
    )


@pytest.mark.parametrize("device", DEVICES)
def test_target_modules_none_uses_all(
    default_vllm_config, dist_init, dummy_model, device
):
    """Test that target_modules=None uses all supported modules."""
    _test_target_modules(
        dummy_model,
        None,
        device,
        expected_lora=[
            ("dense1", ColumnParallelLinearWithLoRA),
            ("layer1.dense1", ColumnParallelLinearWithLoRA),
            ("dense2", RowParallelLinearWithLoRA),
            ("layer1.dense2", RowParallelLinearWithLoRA),
        ],
        expected_no_lora=[],
        vllm_config=default_vllm_config,
    )


@pytest.mark.parametrize("device", DEVICES)
def test_target_modules_match_packed_runtime_modules(
    default_vllm_config, dist_init, dummy_model_gate_up, device
):
    """Packed runtime modules should be selected by their adapter-visible names."""
    _test_target_modules(
        dummy_model_gate_up,
        ["gate_proj"],
        device,
        expected_lora=[("gate_up_proj", MergedColumnParallelLinearWithLoRA)],
        expected_no_lora=[
            ("dense1", ColumnParallelLinearWithLoRA),
            ("dense2", RowParallelLinearWithLoRA),
            ("layer1.dense1", ColumnParallelLinearWithLoRA),
            ("layer1.dense2", RowParallelLinearWithLoRA),
        ],
        vllm_config=default_vllm_config,
    )
