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

import os
from typing import Any, Literal

from pydantic import Field, model_validator
from typing_extensions import Self

from vllm.config.utils import config
from vllm.logger import init_logger
from vllm.utils.hashing import safe_hash

logger = init_logger(__name__)

ProfilerKind = Literal["torch", "cuda", "proton"]
ProtonBackend = Literal["cupti"]
ProtonContext = Literal["shadow", "python"]
ProtonData = Literal["tree", "trace"]
ProtonHook = Literal["triton"]
ProtonOutputFormat = Literal["hatchet", "hatchet_msgpack", "chrome_trace"]


def _is_uri_path(path: str) -> bool:
    """Check if path is a URI (scheme://...), excluding Windows drive letters.

    Supports custom URI schemes like gs://, s3://, hdfs://, etc.
    These paths should not be converted to absolute paths.
    """
    if "://" in path:
        scheme = path.split("://")[0]
        # Windows drive letters are single characters (e.g., C://)
        # Valid URI schemes have more than one character
        return len(scheme) > 1
    return False


@config
class ProfilerConfig:
    """Dataclass which contains profiler config for the engine."""

    profiler: ProfilerKind | None = None
    """Which profiler to use. Defaults to None. Options are:

    - 'torch': Use PyTorch profiler.
    - 'cuda': Use CUDA profiler.
    - 'proton': Use Triton Proton profiler."""

    torch_profiler_dir: str = ""
    """Directory to save torch profiler traces. Both AsyncLLM's CPU traces and
    worker's traces (CPU & GPU) will be saved under this directory. Note that
    it must be an absolute path."""

    proton_profiler_dir: str = ""
    """Directory to save Triton Proton profiles. Each worker writes a
    separate rank-qualified file."""

    proton_context: ProtonContext = "shadow"
    """Proton context source. ``shadow`` records explicit scopes with low
    overhead; ``python`` records Python call stacks."""

    proton_data: ProtonData = "tree"
    """Proton output type. ``tree`` produces Hatchet data and ``trace``
    produces a Chrome trace."""

    proton_backend: ProtonBackend | None = None
    """Proton GPU backend. ``None`` lets Proton select CUPTI automatically."""

    proton_mode: str | None = None
    """Optional backend-specific Proton mode string, such as ``pcsampling``."""

    proton_hook: ProtonHook | None = None
    """Optional Proton hook. Use ``triton`` to add Triton launch metadata."""

    proton_output_format: ProtonOutputFormat | None = None
    """Optional format passed to Proton when finalizing a profile. ``None``
    uses the default format for ``proton_data``."""

    torch_profiler_with_stack: bool = True
    """If `True`, enables stack tracing in the torch profiler. Enabled by default
    as it is useful for debugging. Can be disabled via 
    --profiler-config.torch_profiler_with_stack=false CLI flag."""

    torch_profiler_with_flops: bool = False
    """If `True`, enables FLOPS counting in the torch profiler. Disabled by default."""

    torch_profiler_use_gzip: bool = True
    """If `True`, saves torch profiler traces in gzip format. Enabled by default"""

    torch_profiler_dump_cuda_time_total: bool = True
    """If `True`, dumps total CUDA time in torch profiler traces. Enabled by default."""

    torch_profiler_record_shapes: bool = False
    """If `True`, records tensor shapes in the torch profiler. Disabled by default."""

    torch_profiler_with_memory: bool = False
    """If `True`, enables memory profiling in the torch profiler.
    Disabled by default."""

    capture_torch_profiler: bool = False
    """If `True`, enables a torch profiler during CUDA graph capture on rank 0.
    Traces are saved to a `capture_traces` subdirectory under `torch_profiler_dir`.
    Requires `profiler` to be set to 'torch'."""

    detailed_trace_annotation: bool = False
    """If `True`, uses detailed annotations with roofline metrics (sk, sqsq,
    sqsk) in profiler trace events. If `False`, uses simple annotations with
    only context/generation request counts and token counts.
    Disabled by default."""

    ignore_frontend: bool = False
    """If `True`, disables the front-end profiling of AsyncLLM when using the
    'torch' profiler. This is needed to reduce overhead when using delay/limit options,
    since the front-end profiling does not track iterations and will capture the
    entire range.
    """

    delay_iterations: int = Field(default=0, ge=0)
    """Number of engine iterations to skip before starting profiling.
    Defaults to 0, meaning profiling starts immediately after receiving /start_profile.
    """

    max_iterations: int = Field(default=0, ge=0)
    """Maximum number of engine iterations to profile after starting profiling.
    Defaults to 0, meaning no limit.
    """

    warmup_iterations: int = Field(default=0, ge=0)
    """Number of warmup iterations for PyTorch profiler schedule.
    During warmup, the profiler runs but data is discarded. This helps reduce
    noise from JIT compilation and other one-time costs in the profiled trace.
    Defaults to 0 (schedule-based profiling disabled, recording all iterations).
    Set to a positive value (e.g., 2) to enable schedule-based profiling.
    """

    active_iterations: int = Field(default=5, ge=1)
    """Number of active iterations for PyTorch profiler schedule.
    This is the number of iterations where profiling data is actually collected.
    Defaults to 5 active iterations.
    """

    wait_iterations: int = Field(default=0, ge=0)
    """Number of wait iterations for PyTorch profiler schedule.
    During wait, the profiler is completely off with zero overhead.
    This allows skipping initial iterations before warmup begins.
    Defaults to 0 (no wait period).
    """

    def compute_hash(self) -> str:
        """
        WARNING: Whenever a new field is added to this config,
        ensure that it is included in the factors list if
        it affects the computation graph.

        Provide a hash that uniquely identifies all the configs
        that affect the structure of the computation
        graph from input ids/embeddings to the final hidden states,
        excluding anything before input ids/embeddings and after
        the final hidden states.
        """
        # no factors to consider.
        # this config will not affect the computation graph.
        factors: list[Any] = []
        hash_str = safe_hash(str(factors).encode(), usedforsecurity=False).hexdigest()
        return hash_str

    @model_validator(mode="after")
    def _validate_profiler_config(self) -> Self:
        has_delay_or_limit = self.delay_iterations > 0 or self.max_iterations > 0
        if self.profiler == "torch" and has_delay_or_limit and not self.ignore_frontend:
            logger.warning_once(
                "Using 'torch' profiler with delay_iterations or max_iterations "
                "while ignore_frontend is False may result in high overhead."
            )

        torch_profiler_dir = self.torch_profiler_dir
        if torch_profiler_dir and self.profiler != "torch":
            raise ValueError(
                "torch_profiler_dir is only applicable when profiler is set to 'torch'"
            )
        if self.profiler == "torch" and not torch_profiler_dir:
            raise ValueError("torch_profiler_dir must be set when profiler is 'torch'")

        # Support any URI scheme (gs://, s3://, hdfs://, etc.)
        # These paths should not be converted to absolute paths
        if torch_profiler_dir and not _is_uri_path(torch_profiler_dir):
            self.torch_profiler_dir = os.path.abspath(
                os.path.expanduser(torch_profiler_dir)
            )

        proton_profiler_dir = self.proton_profiler_dir
        non_default_proton_options = [
            name
            for name, value, default in (
                ("proton_profiler_dir", proton_profiler_dir, ""),
                ("proton_context", self.proton_context, "shadow"),
                ("proton_data", self.proton_data, "tree"),
                ("proton_backend", self.proton_backend, None),
                ("proton_mode", self.proton_mode, None),
                ("proton_hook", self.proton_hook, None),
                ("proton_output_format", self.proton_output_format, None),
            )
            if value != default
        ]
        if self.profiler != "proton" and non_default_proton_options:
            options = ", ".join(non_default_proton_options)
            raise ValueError(
                f"{options} only applicable when profiler is set to 'proton'"
            )
        if self.profiler == "proton" and not proton_profiler_dir:
            raise ValueError(
                "proton_profiler_dir must be set when profiler is 'proton'"
            )
        if proton_profiler_dir:
            if _is_uri_path(proton_profiler_dir):
                raise ValueError("proton_profiler_dir must be a local directory")
            self.proton_profiler_dir = os.path.abspath(
                os.path.expanduser(proton_profiler_dir)
            )

        if self.profiler == "proton":
            output_format = self.proton_output_format
            if output_format == "chrome_trace" and self.proton_data != "trace":
                raise ValueError("chrome_trace output requires proton_data='trace'")
            if (
                output_format in ("hatchet", "hatchet_msgpack")
                and self.proton_data != "tree"
            ):
                raise ValueError(f"{output_format} output requires proton_data='tree'")

        if self.capture_torch_profiler and self.profiler != "torch":
            raise ValueError(
                "capture_torch_profiler is only applicable when profiler is "
                "set to 'torch'"
            )

        return self