# SPDX-FileCopyrightText: NVIDIA CORPORATION & AFFILIATES
# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: Apache-2.0
#
# 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.

from __future__ import annotations

import json
import logging
import re
from functools import cache
from pathlib import Path
from typing import Optional, cast

from pydantic import ConfigDict, Field, field_validator

from cloudai.core import GitRepo, Installable, JobStatusResult, System, TestRun
from cloudai.models.workload import CmdArgs
from cloudai.workloads.common.llm_serving import (
    CustomBash,
    LLMServingArgs,
    LLMServingCmdArgs,
    LLMServingTestDefinition,
    all_gpu_ids,
    calculate_decode_gpu_ids,
    calculate_prefill_gpu_ids,
    validate_custom_bash_patterns,
)

VLLM_SERVE_LOG_FILE = "vllm-serve.log"
VLLM_BENCH_LOG_FILE = "vllm-bench.log"
VLLM_BENCH_JSON_FILE = "vllm-bench.json"
VLLM_GSM8K_JSON_FILE = "vllm-gsm8k.json"
VLLM_SEMANTIC_EVAL_LOG_FILE = "vllm-semantic-eval.log"



[docs]
class VllmArgs(LLMServingArgs):
    """Base command arguments for vLLM instances."""

    ray_head: VllmRayStartArgs | None = Field(
        default=None,
        description="Arguments appended to the Ray head startup command for multi-node vLLM roles.",
    )
    ray_worker: VllmRayStartArgs | None = Field(
        default=None,
        description="Arguments appended to the Ray worker startup command for multi-node vLLM roles.",
    )
    nixl_threads: int | list[int] | None = Field(
        default=None,
        description="Set ``kv_connector_extra_config.num_threads`` for ``--kv-transfer-config`` CLI argument.",
    )

    @property
    def serve_args_exclude(self) -> set[str]:
        return super().serve_args_exclude | {"nixl_threads", "ray_head", "ray_worker"}


[docs]
    def serialize_serve_arg(self, key: str, value: object) -> list[str]:
        opt = f"--{key.replace('_', '-')}"
        if isinstance(value, bool):
            return [opt] if value else [f"--no-{key.replace('_', '-')}"]
        return super().serialize_serve_arg(key, value)
class VllmRayStartArgs(CmdArgs):
    """Ray startup arguments for vLLM multi-node serving roles."""

    model_config = ConfigDict(extra="allow")

    head: bool | list[bool] | None = Field(default=None, description="Emit ``--head`` for Ray head startup.")
    port: int | str | list[int] | list[str] | None = Field(default=None, description="Ray head port.")
    address: str | list[str] | None = Field(default=None, description="Ray head address for worker startup.")
    block: bool | list[bool] | None = Field(default=None, description="Emit ``--block`` for Ray worker startup.")







[docs]
class VllmCmdArgs(LLMServingCmdArgs[VllmArgs]):
    """vLLM serve command arguments."""

    model_config = ConfigDict(extra="forbid")  # arbitrary fields are allowed per decode/prefill, not here

    proxy_script: str = "/opt/vllm/tests/v1/kv_connector/nixl_integration/toy_proxy_server.py"
    healthcheck: str = Field(default="/healthcheck", description="vLLM server healthcheck endpoint.")
    proxy_healthcheck: str = Field(
        default="/healthcheck",
        description="vLLM disaggregated proxy/router healthcheck endpoint.",
    )

    model: str = "Qwen/Qwen3-0.6B"
    prefill: VllmArgs | None = Field(
        default=None,
        description="Prefill instance arguments. If not set, a single instance without disaggregation will be used.",
    )
    decode: VllmArgs = Field(default_factory=VllmArgs, description="Decode instance arguments.")




[docs]
class VllmBenchCmdArgs(CmdArgs):
    """vLLM bench serve command arguments."""

    random_input_len: int = 16
    random_output_len: int = 128
    max_concurrency: int = 16
    num_prompts: int = 30




[docs]
class VllmSemanticEvalCmdArgs(CmdArgs):
    """vLLM semantic validation command arguments."""

    model_config = ConfigDict(extra="forbid")

    entrypoint: str = "python3 /opt/vllm/tests/evals/gsm8k/gsm8k_eval.py"
    cli: str = "--host {host} --port {port} --num-questions 200 --save-results {output_path}/vllm-gsm8k.json"




[docs]
class VllmTestDefinition(LLMServingTestDefinition[VllmCmdArgs]):
    """Test object for vLLM."""

    bench_cmd_args: VllmBenchCmdArgs = VllmBenchCmdArgs()
    semantic_eval_cmd_args: VllmSemanticEvalCmdArgs | None = None
    proxy_script_repo: GitRepo | None = None
    custom_bash: CustomBash | None = None

    @field_validator("custom_bash", mode="after")
    @classmethod
    def validate_custom_bash(cls, custom_bash: CustomBash | None) -> CustomBash | None:
        return validate_custom_bash_patterns(custom_bash)

    @property
    def extra_installables(self) -> list[Installable]:
        installables: list[Installable] = []
        if self.proxy_script_repo:
            installables.append(self.proxy_script_repo)
        return installables

    @staticmethod
    def _validate_vllm_parallelism_constraints(role: str, args: VllmArgs, gpu_count: int) -> bool:
        tp = cast(int, getattr(args, "tensor_parallel_size", 1))
        pp = cast(int, getattr(args, "pipeline_parallel_size", 1))
        dp = cast(int, getattr(args, "data_parallel_size", 1))
        ep_enabled = cast(bool, getattr(args, "enable_expert_parallel", False))
        all2all_backend = cast(str, getattr(args, "all2all_backend", ""))

        constraint1 = (tp * pp * dp) <= gpu_count
        if not constraint1:
            logging.error(
                "vLLM %s constraint failed: (tp * pp * dp) <= num_gpus. tp=%s pp=%s dp=%s num_gpus=%s",
                role,
                tp,
                pp,
                dp,
                gpu_count,
            )
            return False

        using_flashinfer_all2allv = all2all_backend == "flashinfer_all2allv"
        constraint2 = not (using_flashinfer_all2allv and dp > 1 and ep_enabled)
        if not constraint2:
            logging.error(
                "vLLM %s constraint failed: flashinfer_all2allv only works with DP=1, or with DP>1 and expert "
                "parallel disabled. all2all_backend=%s dp=%s expert_parallel=%s",
                role,
                all2all_backend,
                dp,
                ep_enabled,
            )
            return False

        return True

    def constraint_check(self, tr: TestRun, system: Optional[System]) -> bool:
        system_gpus_per_node = getattr(system, "gpus_per_node", None) if system is not None else None
        num_nodes = tr.nnodes
        local_gpu_count = len(all_gpu_ids(self, system_gpus_per_node))

        if self.cmd_args.prefill is None:
            return self._validate_vllm_parallelism_constraints(
                role="decode",
                args=self.cmd_args.decode,
                gpu_count=local_gpu_count * num_nodes,
            )

        prefill_nodes_value = self.cmd_args.prefill.num_nodes
        decode_nodes_value = self.cmd_args.decode.num_nodes
        if prefill_nodes_value is None and decode_nodes_value is None:
            if num_nodes > 2:
                logging.error(
                    "vLLM disaggregated mode over more than 2 nodes requires both prefill.num_nodes and "
                    "decode.num_nodes."
                )
                return False
            prefill_nodes = 1
            decode_nodes = 1
        elif not isinstance(prefill_nodes_value, int) or not isinstance(decode_nodes_value, int):
            logging.error("vLLM disaggregated role node counts must both be single integers or both be omitted.")
            return False
        elif prefill_nodes_value <= 0 or decode_nodes_value <= 0:
            logging.error(
                "vLLM disaggregated role node counts must be positive integers. prefill=%s decode=%s",
                prefill_nodes_value,
                decode_nodes_value,
            )
            return False
        elif num_nodes == 1 and prefill_nodes_value == 1 and decode_nodes_value == 1:
            prefill_nodes = 1
            decode_nodes = 1
        elif prefill_nodes_value + decode_nodes_value != num_nodes:
            logging.error(
                "vLLM disaggregated role node counts must sum to allocated nodes. prefill=%s decode=%s allocated=%s",
                prefill_nodes_value,
                decode_nodes_value,
                num_nodes,
            )
            return False
        else:
            prefill_nodes = prefill_nodes_value
            decode_nodes = decode_nodes_value

        return self._validate_vllm_parallelism_constraints(
            role="prefill",
            args=self.cmd_args.prefill,
            gpu_count=len(calculate_prefill_gpu_ids(self, num_nodes, system_gpus_per_node)) * prefill_nodes,
        ) and self._validate_vllm_parallelism_constraints(
            role="decode",
            args=self.cmd_args.decode,
            gpu_count=len(calculate_decode_gpu_ids(self, num_nodes, system_gpus_per_node)) * decode_nodes,
        )

    def was_run_successful(self, tr: TestRun) -> JobStatusResult:
        log_path = tr.output_path / VLLM_BENCH_LOG_FILE
        if not log_path.is_file():
            return JobStatusResult(is_successful=False, error_message=f"vLLM bench log not found in {tr.output_path}.")

        has_results_marker = False
        with log_path.open("r") as f:
            for line in f:
                if "============ Serving Benchmark Result ============" in line:
                    has_results_marker = True
                    continue
                if has_results_marker and "Successful requests:" in line:
                    try:
                        num_successful_requests = int(line.split()[2])
                        if num_successful_requests > 0:
                            if self.semantic_eval_cmd_args is not None:
                                accuracy = parse_vllm_semantic_accuracy(tr.output_path)
                                if accuracy is None:
                                    return JobStatusResult(
                                        is_successful=False,
                                        error_message=f"vLLM semantic accuracy not found in {tr.output_path}.",
                                    )
                            return JobStatusResult(is_successful=True)
                    except Exception as e:
                        logging.debug(f"Error parsing number of successful requests: {e}")

        return JobStatusResult(
            is_successful=False, error_message=f"vLLM bench log does not contain benchmark result in {tr.output_path}."
        )
@cache
def parse_vllm_semantic_accuracy(output_path: Path) -> float | None:
    """Parse vLLM semantic validation accuracy from JSON results or the eval log."""
    json_path = output_path / VLLM_GSM8K_JSON_FILE
    if json_path.is_file():
        try:
            data = json.loads(json_path.read_text(encoding="utf-8"))
            accuracy = data.get("accuracy") if isinstance(data, dict) else None
            if isinstance(accuracy, (int, float)):
                return float(accuracy)
        except Exception as e:
            logging.debug(f"Error parsing vLLM semantic JSON output: {e}")

    log_path = output_path / VLLM_SEMANTIC_EVAL_LOG_FILE
    if not log_path.is_file():
        return None

    pattern = re.compile(r"\bAccuracy:\s*([0-9]*\.?[0-9]+)")
    with log_path.open(encoding="utf-8", errors="ignore") as f:
        for line in f:
            match = pattern.search(line)
            if match:
                return float(match.group(1))

    return None