# SPDX-FileCopyrightText: NVIDIA CORPORATION & AFFILIATES
# Copyright (c) 2025-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 typing import Literal, Optional

from pydantic import Field, model_validator

from cloudai.core import CmdArgs, System, TestRun
from cloudai.workloads.common.nixl import NIXLBaseCmdArgs, NIXLBaseTestDefinition


class MatgenCmdArgs(CmdArgs):
    """Command args for matgen script."""

    ppn: int | None = None



[docs]
class NixlPerftestCmdArgs(NIXLBaseCmdArgs):
    """CmdArgs for NIXL Perftest."""

    subtest: Literal["sequential-ct-perftest"]
    perftest_script: str = "/workspace/nixl/benchmark/kvbench/main.py"
    matgen_script: str = "/workspace/nixl/benchmark/kvbench/test/inference_workload_matgen.py"
    python_executable: str = "python"

    num_user_requests: int | list[int]
    batch_size: int | list[int]
    num_prefill_nodes: int | list[int]
    num_decode_nodes: int | list[int]
    isl_mean: int | list[int] | None = None
    isl_scale: int | list[int] | None = None
    prefill_tp: int | list[int] = 1
    prefill_pp: int | list[int] = 1
    prefill_cp: int | list[int] = 1
    decode_tp: int | list[int] = 1
    decode_pp: int | list[int] = 1
    decode_cp: int | list[int] = 1

    # model or model configuration
    model: str | list[str] | None = None
    hidden_size: int | None = None
    num_layers: int | None = None
    num_heads: int | None = None
    num_kv_heads: int | None = None
    dtype_size: int | None = None

    matgen_args: MatgenCmdArgs = Field(default_factory=MatgenCmdArgs)

    @model_validator(mode="after")
    def model_vs_custom(self):
        if self.model is None and (
            self.hidden_size is None
            or self.num_layers is None
            or self.num_heads is None
            or self.num_kv_heads is None
            or self.dtype_size is None
        ):
            raise ValueError(
                "If 'model' is None, 'hidden_size', 'num_layers', 'num_heads', 'num_kv_heads', and 'dtype_size' "
                "must be specified."
            )

        if self.model is not None and (
            self.hidden_size is not None
            or self.num_layers is not None
            or self.num_heads is not None
            or self.num_kv_heads is not None
            or self.dtype_size is not None
        ):
            raise ValueError(
                "If 'model' is specified, 'hidden_size', 'num_layers', 'num_heads', 'num_kv_heads', and 'dtype_size' "
                "must be None."
            )

        return self




[docs]
class NixlPerftestTestDefinition(NIXLBaseTestDefinition[NixlPerftestCmdArgs]):
    """TestDefinition for NixlPerftest."""

    def constraint_check(self, tr: TestRun, system: Optional[System]) -> bool:
        decode_tp = int(tr.test.cmd_args.decode_tp)
        decode_nodes = int(tr.test.cmd_args.num_decode_nodes)
        prefill_tp = int(tr.test.cmd_args.prefill_tp)
        prefill_nodes = int(tr.test.cmd_args.num_prefill_nodes)

        return decode_tp / decode_nodes == prefill_tp / prefill_nodes