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

import contextlib
import inspect
import os
import shutil
from functools import wraps
from types import FrameType

import psutil
import torch
from compressed_tensors.distributed import is_distributed, is_source_process
from compressed_tensors.offload.convert import from_accelerate
from loguru import logger
from transformers import PreTrainedModel
from transformers.models.auto.modeling_auto import _BaseAutoModelClass


__all__ = ["load_offloaded_model"]


cls_to_patch = _BaseAutoModelClass | PreTrainedModel


@contextlib.contextmanager
def load_offloaded_model(extra_cpu_mem: int = 5e9):
    """
    Context manager used to load a transformers model with offloading implemented by
    compressed-tensors.

    The model is first loaded with accelerate's offloading, then convereted into
    offloading implemented by compressed-tensors. If a distributed environment has been
    initialized, then rank 0 loads the weights while other ranks load on the meta
    device, then the offload is shared across ranks during conversion.

    In addition to the standard `device_map` options, this context also supports
    `device_map="auto_offload"`, which means that the model will load as many parameters
    can fit onto the cpu, and any extra parameters will be loaded on disk.

    :param extra_cpu_mem: extra cpu memory to reserve for any operations not related to
        model loading (bytes). Defaults to 5Gb.
    """
    frame = _get_caller_frame()

    with contextlib.ExitStack() as stack:
        for obj in frame.f_globals.values():
            if isinstance(obj, type) and issubclass(obj, cls_to_patch):
                stack.enter_context(patch_from_pretrained(obj, extra_cpu_mem))

        yield


@contextlib.contextmanager
def patch_from_pretrained(obj: cls_to_patch, extra_cpu_mem: int):
    original_func = obj.from_pretrained.__func__

    @wraps(original_func)
    def from_pretrained(cls, *args, **kwargs):
        kwargs.setdefault("device_map", None)

        # Rank 0 does loading, other ranks init on meta device
        if not is_source_process():
            kwargs["device_map"] = "meta"
            # Workaround: transformers v5 tie_weights() calls torch.equal() on
            # meta tensors which is unsupported. Since rank 0 broadcasts the real
            # weights, we can safely skip tying on non-rank workers.
            kwargs.setdefault("tie_word_embeddings", False)

        # Intercept `auto_offload`: same as "auto", but only cpu/disk are visible
        elif kwargs["device_map"] == "auto_offload":
            kwargs["device_map"] = "auto"
            if "max_memory" not in kwargs:
                kwargs["max_memory"] = _get_cpu_memory(extra_cpu_mem)

        # Unless the user specifies, use our memory estimates, which take into
        # account distributed setups and extra cpu reserved memory
        elif "max_memory" not in kwargs:
            kwargs["max_memory"] = _get_device_memory() | _get_cpu_memory(extra_cpu_mem)

        model = original_func(cls, *args, **kwargs)
        from_accelerate(model)  # rank 0 shares weights with ranks via offload/broadcast
        return model

    try:
        obj.from_pretrained = from_pretrained.__get__(obj)
        yield
    finally:
        obj.from_pretrained = original_func.__get__(obj)


def _get_device_memory() -> dict[int, int]:
    if is_distributed():
        index = torch.accelerator.current_device_index()
        return {index: torch.accelerator.get_memory_info(index)[1]}
    else:
        return {
            index: torch.accelerator.get_memory_info(index)[1]
            for index in range(torch.accelerator.device_count())
        }


def _get_cpu_memory(extra_cpu_mem: int) -> dict[str, int]:
    if is_distributed():
        return {"cpu": _get_shared_memory() - extra_cpu_mem}
    else:
        return {"cpu": psutil.virtual_memory().available - extra_cpu_mem}


def _get_shared_memory() -> int:
    linux_shm_path = "/dev/shm"
    if os.path.exists(linux_shm_path):
        total, _used, _free = shutil.disk_usage(linux_shm_path)
        return total

    else:
        logger.warning(
            "Could not find shared memory at `/dev/shm`. Please add platform suppport"
        )
        return psutil.virtual_memory().available


def _get_caller_frame() -> FrameType:
    frame = inspect.currentframe()
    frame = frame.f_back.f_back  # skip this function's caller's frame
    while frame is not None and "contextlib" in frame.f_code.co_filename:
        frame = frame.f_back  # skip contextlib frames

    if frame is None:
        raise RuntimeError("Could not find caller frame")

    return frame
