schema.py

#!/usr/bin/env python3
# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: Apache-2.0

"""
Frozen dataclass schema definitions for job configuration.

Uses marshmallow_dataclass for type-safe configuration with validation.
All config classes are frozen (immutable) after creation.

Backend configs are defined in srtctl.backends.configs/ for modularity.
"""

import builtins
import itertools
import logging
from collections.abc import Iterator, Mapping
from dataclasses import field
from enum import Enum
from pathlib import Path
from typing import (
    TYPE_CHECKING,
    Annotated,
    Any,
    ClassVar,
    Literal,
)

import yaml
from marshmallow import Schema, ValidationError, fields
from marshmallow_dataclass import dataclass

from srtctl.backends import (
    BackendConfig,
    MockerProtocol,
    SGLangProtocol,
    TRTLLMProtocol,
    VLLMProtocol,
)
from srtctl.core.formatting import (
    FormattablePath,
    FormattablePathField,
)

if TYPE_CHECKING:
    pass

logger = logging.getLogger(__name__)


# ============================================================================
# Reporting Configuration
# ============================================================================


@dataclass(frozen=True)
class ReportingStatusConfig:
    """Status reporting configuration."""

    endpoint: str | None = None
    endpoints: list[str] | None = None

    Schema: ClassVar[type[Schema]] = Schema


@dataclass(frozen=True)
class ReportingConfig:
    """Reporting configuration for status updates, AI analysis, and log exports."""

    status: ReportingStatusConfig | None = None
    ai_analysis: "AIAnalysisConfig | None" = None
    s3: "S3Config | None" = None

    Schema: ClassVar[type[Schema]] = Schema


# ============================================================================
# Cluster Configuration (srtslurm.yaml)
# ============================================================================


# Default prompt template for AI-powered failure analysis
DEFAULT_AI_ANALYSIS_PROMPT = """
You are analyzing benchmark failure logs for an LLM serving system (SGLang/Dynamo).

You have access to:
- Log files in {log_dir}
- The `gh` CLI tool (authenticated) to search GitHub PRs

Your task:
1. Read the log files and identify the root cause of failure
2. Search recent PRs (last {pr_days} days) in {repos} for potentially related changes
3. Write your analysis to ai_analysis.md in {log_dir}

Your analysis should include:
- Summary of the failure
- Root cause identification
- Key error messages found
- Related PRs (if any)
- Suggested next steps

Start by listing and reading the log files, then investigate.
"""


@dataclass(frozen=True)
class AIAnalysisConfig:
    """AI-powered failure analysis configuration.

    This config is typically set in srtslurm.yaml (cluster config) to centralize
    secrets and allow cluster-wide customization. Individual job configs can
    override with `ai_analysis.enabled: false` to disable for specific jobs.

    Uses OpenRouter for Claude Code authentication, which provides a simple API key
    approach that works well in headless/automated environments.
    See: https://openrouter.ai/docs/guides/claude-code-integration

    Attributes:
        enabled: Whether to run AI analysis on benchmark failures
        openrouter_api_key: OpenRouter API key (falls back to OPENROUTER_API_KEY env var)
        gh_token: GitHub token for gh CLI (falls back to GH_TOKEN env var)
        repos_to_search: GitHub repos to search for related PRs
        pr_search_days: Number of days to look back for PRs
        prompt: Custom prompt template (uses DEFAULT_AI_ANALYSIS_PROMPT if None)
            Available variables: {log_dir}, {repos}, {pr_days}
    """

    enabled: bool = False
    openrouter_api_key: str | None = None
    gh_token: str | None = None
    repos_to_search: list[str] = field(default_factory=lambda: ["sgl-project/sglang", "ai-dynamo/dynamo"])
    pr_search_days: int = 14
    prompt: str | None = None

    def get_prompt(self, log_dir: str) -> str:
        """Get the formatted prompt for AI analysis.

        Args:
            log_dir: Path to the log directory

        Returns:
            Formatted prompt string
        """
        template = self.prompt or DEFAULT_AI_ANALYSIS_PROMPT
        repos_str = ", ".join(self.repos_to_search)
        return template.format(
            log_dir=log_dir,
            repos=repos_str,
            pr_days=self.pr_search_days,
        )

    Schema: ClassVar[type[Schema]] = Schema


@dataclass(frozen=True)
class S3Config:
    """S3 upload configuration for log artifacts.

    Attributes:
        bucket: S3 bucket name
        prefix: Optional prefix/path within bucket (e.g., "srtslurm/logs")
        region: AWS region (e.g., "us-west-2")
        endpoint_url: Custom S3-compatible endpoint URL (optional)
        access_key_id: AWS access key ID (falls back to AWS_ACCESS_KEY_ID env var)
        secret_access_key: AWS secret access key (falls back to AWS_SECRET_ACCESS_KEY env var)
    """

    bucket: str
    prefix: str | None = None
    region: str | None = None
    endpoint_url: str | None = None
    access_key_id: str | None = None
    secret_access_key: str | None = None

    Schema: ClassVar[type[Schema]] = Schema


@dataclass
class ClusterConfig:
    """Cluster configuration from srtslurm.yaml."""

    cluster: str | None = None  # Cluster name for status reporting
    default_account: str | None = None
    default_partition: str | None = None
    default_time_limit: str | None = None
    gpus_per_node: int | None = None
    network_interface: str | None = None
    use_gpus_per_node_directive: bool = True
    use_segment_sbatch_directive: bool = True
    use_exclusive_sbatch_directive: bool = False
    srtctl_root: str | None = None
    output_dir: str | None = None  # Custom output directory for job logs
    model_paths: dict[str, str] | None = None
    containers: dict[str, str] | None = None
    cloud: dict[str, str] | None = None
    # Cluster-level container mounts (host_path -> container_path)
    # Applied to all jobs on this cluster, useful for cluster-specific paths
    default_mounts: dict[str, str] | None = None
    reporting: ReportingConfig | None = None

    Schema: ClassVar[type[Schema]] = Schema


# ============================================================================
# Enums
# ============================================================================


class GpuType(str, Enum):
    GB200 = "gb200"
    GB300 = "gb300"
    H100 = "h100"


class Precision(str, Enum):
    FP4 = "fp4"
    FP8 = "fp8"
    FP16 = "fp16"
    BF16 = "bf16"


class BenchmarkType(str, Enum):
    MANUAL = "manual"
    CUSTOM = "custom"
    SA_BENCH = "sa-bench"
    ROUTER = "router"
    MOONCAKE_ROUTER = "mooncake-router"
    TRACE_REPLAY = "trace-replay"
    AIME = "aime"
    MMLU = "mmlu"
    GPQA = "gpqa"
    GSM8K = "gsm8k"
    LONGBENCHV2 = "longbenchv2"


class ProfilingType(str, Enum):
    NSYS = "nsys"
    TORCH = "torch"
    NONE = "none"


class TelemetryProvider(str, Enum):
    SCRAPER = "scraper"


# ============================================================================
# Marshmallow Custom Fields
# ============================================================================


class BackendConfigField(fields.Field):
    """Marshmallow field for polymorphic backend deserialization based on type."""

    def _deserialize(
        self,
        value: Any,
        attr: str | None,
        data: Mapping[str, Any] | None,
        **kwargs,
    ) -> BackendConfig:
        """Deserialize backend config based on 'type' field."""
        if value is None:
            # Default to SGLang
            return SGLangProtocol()

        if isinstance(value, SGLangProtocol | TRTLLMProtocol | VLLMProtocol | MockerProtocol):
            return value

        if not isinstance(value, dict):
            raise ValidationError(f"Expected dict for backend config, got {type(value).__name__}")

        # Get backend type from the value dict
        backend_type = value.get("type", "sglang")

        if backend_type == "sglang":
            schema = SGLangProtocol.Schema()
            return schema.load(value)
        elif backend_type == "trtllm":
            schema = TRTLLMProtocol.Schema()
            return schema.load(value)
        elif backend_type == "vllm":
            schema = VLLMProtocol.Schema()
            return schema.load(value)
        elif backend_type == "mocker":
            schema = MockerProtocol.Schema()
            return schema.load(value)
        else:
            raise ValidationError(
                f"Unknown backend type: {backend_type!r}. Supported types: sglang, trtllm, vllm, mocker"
            )

    def _serialize(self, value: Any | None, attr: str | None, obj: Any, **kwargs) -> Any:
        """Serialize backend config to dict."""
        if value is None:
            return None
        if isinstance(value, SGLangProtocol):
            return SGLangProtocol.Schema().dump(value)
        if isinstance(value, TRTLLMProtocol):
            return TRTLLMProtocol.Schema().dump(value)
        if isinstance(value, VLLMProtocol):
            return VLLMProtocol.Schema().dump(value)
        if isinstance(value, MockerProtocol):
            return MockerProtocol.Schema().dump(value)
        return value


class SweepConfigField(fields.Field):
    """Marshmallow field for SweepConfig."""

    def _deserialize(self, value: Any, attr: str | None, data: Mapping[str, Any] | None, **kwargs) -> Any:
        if value is None:
            return None
        if isinstance(value, SweepConfig):
            return value
        if not isinstance(value, dict):
            raise ValidationError(f"Expected dict for sweep config, got {type(value).__name__}")

        mode = value.get("mode", "zip")
        parameters: dict[str, list[Any]] = {}

        if "parameters" in value:
            for key, val in value["parameters"].items():
                if not isinstance(val, list):
                    raise ValidationError(f"Sweep parameter '{key}' must be a list")
                parameters[key] = val
        else:
            for key, val in value.items():
                if key == "mode":
                    continue
                if not isinstance(val, list):
                    raise ValidationError(f"Sweep parameter '{key}' must be a list")
                parameters[key] = val

        return SweepConfig(mode=mode, parameters=parameters)

    def _serialize(self, value: Any | None, attr: str | None, obj: Any, **kwargs) -> Any:
        if value is None:
            return None
        if isinstance(value, SweepConfig):
            result: dict[str, Any] = {"mode": value.mode}
            result.update(value.parameters)
            return result
        return value


# ============================================================================
# Sub-Configuration Dataclasses (all frozen)
# ============================================================================


@dataclass(frozen=True)
class SweepConfig:
    """Configuration for benchmark parameter sweeps."""

    mode: Literal["zip", "grid"] = "zip"
    parameters: dict[str, list[Any]] = field(default_factory=dict)

    def get_combinations(self) -> Iterator[dict[str, Any]]:
        if not self.parameters:
            yield {}
            return

        if self.mode == "zip":
            param_names = list(self.parameters.keys())
            param_lists = [self.parameters[name] for name in param_names]
            for values in zip(*param_lists, strict=False):
                yield dict(zip(param_names, values, strict=False))
        else:
            param_names = list(self.parameters.keys())
            param_lists = [self.parameters[name] for name in param_names]
            for values in itertools.product(*param_lists):
                yield dict(zip(param_names, values, strict=False))

    def __len__(self) -> int:
        if not self.parameters:
            return 1
        if self.mode == "zip":
            return len(next(iter(self.parameters.values())))
        result = 1
        for param_list in self.parameters.values():
            result *= len(param_list)
        return result

    Schema: ClassVar[type[Schema]] = Schema


@dataclass(frozen=True)
class ModelConfig:
    """Model configuration."""

    path: str
    container: str
    precision: str

    Schema: ClassVar[type[Schema]] = Schema


@dataclass(frozen=True)
class IdentityModelConfig:
    """Virtual model identity for runtime verification."""

    repo: str | None = None  # HuggingFace model ID, e.g. "nvidia/Kimi-K2.5-NVFP4"
    revision: str | None = None  # HuggingFace git commit SHA

    Schema: ClassVar[type[Schema]] = Schema


@dataclass(frozen=True)
class IdentityContainerConfig:
    """Container identity for reproduction (not verified at runtime).

    Recorded so others can pull the same container image to reproduce.
    Cannot be verified at runtime — Pyxis/enroot strips provenance during import.
    """

    image: str | None = None  # Docker URI, e.g. "gitlab-master:5005/.../trtllm-arm64"

    Schema: ClassVar[type[Schema]] = Schema


@dataclass(frozen=True)
class IdentityConfig:
    """Virtual identity for runtime verification and reproduction.

    These fields declare what *should* be running. They are not used for
    launching — only for verifying the runtime fingerprint matches expectations
    and for helping others reproduce the run.

    - model: HF repo + revision (verified against download metadata at runtime)
    - container: Docker image URI (recorded for reproduction, not verified)
    - frameworks: expected versions for dynamo + one engine (verified via importlib.metadata)
    """

    model: IdentityModelConfig = field(default_factory=IdentityModelConfig)
    container: IdentityContainerConfig = field(default_factory=IdentityContainerConfig)
    frameworks: dict[str, str] = field(default_factory=dict)

    Schema: ClassVar[type[Schema]] = Schema


@dataclass(frozen=True)
class ResourceConfig:
    """Resource allocation configuration."""

    gpu_type: str
    gpus_per_node: int = 4

    # Disaggregated mode
    prefill_nodes: int | None = None
    decode_nodes: int | None = None
    prefill_workers: int | None = None
    decode_workers: int | None = None

    # Aggregated mode
    agg_nodes: int | None = None
    agg_workers: int | None = None

    # Explicit GPUs per worker (override computed values)
    # Use data_key to map from YAML field names to internal attribute names
    _explicit_gpus_per_prefill: int | None = field(
        default=None,
        metadata={
            "marshmallow_field": fields.Integer(
                data_key="gpus_per_prefill",
                load_default=None,
                allow_none=True,
            )
        },
    )
    _explicit_gpus_per_decode: int | None = field(
        default=None,
        metadata={
            "marshmallow_field": fields.Integer(
                data_key="gpus_per_decode",
                load_default=None,
                allow_none=True,
            )
        },
    )
    _explicit_gpus_per_agg: int | None = field(
        default=None,
        metadata={
            "marshmallow_field": fields.Integer(
                data_key="gpus_per_agg",
                load_default=None,
                allow_none=True,
            )
        },
    )

    @property
    def is_disaggregated(self) -> bool:
        return self.prefill_nodes is not None or self.decode_nodes is not None

    @property
    def total_nodes(self) -> int:
        if self.is_disaggregated:
            return (self.prefill_nodes or 0) + (self.decode_nodes or 0)
        return self.agg_nodes or 1

    @property
    def num_prefill(self) -> int:
        return self.prefill_workers or 0

    @property
    def num_decode(self) -> int:
        return self.decode_workers or 0

    @property
    def num_agg(self) -> int:
        return self.agg_workers or 0

    @property
    def gpus_per_prefill(self) -> int:
        # Use explicit value if set
        if self._explicit_gpus_per_prefill is not None:
            return self._explicit_gpus_per_prefill
        # Fall back to computed value
        if self.prefill_nodes and self.prefill_workers:
            return (self.prefill_nodes * self.gpus_per_node) // self.prefill_workers
        return self.gpus_per_node

    @property
    def gpus_per_decode(self) -> int:
        # Use explicit value if set
        if self._explicit_gpus_per_decode is not None:
            return self._explicit_gpus_per_decode
        # Fall back to computed value
        if self.decode_nodes and self.decode_workers:
            return (self.decode_nodes * self.gpus_per_node) // self.decode_workers
        # decode_nodes=0 with decode_workers means "share nodes with prefill"
        # Inherit TP from prefill in this case
        if self.decode_nodes == 0 and self.decode_workers:
            return self.gpus_per_prefill
        return self.gpus_per_node

    @property
    def gpus_per_agg(self) -> int:
        # Use explicit value if set
        if self._explicit_gpus_per_agg is not None:
            return self._explicit_gpus_per_agg
        # Fall back to computed value
        if self.agg_nodes and self.agg_workers:
            return (self.agg_nodes * self.gpus_per_node) // self.agg_workers
        return self.gpus_per_node

    @property
    def prefill_gpus(self) -> int:
        """Total GPUs used by all prefill workers."""
        return self.num_prefill * self.gpus_per_prefill

    @property
    def decode_gpus(self) -> int:
        """Total GPUs used by all decode workers."""
        return self.num_decode * self.gpus_per_decode

    Schema: ClassVar[type[Schema]] = Schema


@dataclass(frozen=True)
class SlurmConfig:
    """SLURM job settings."""

    account: str | None = None
    partition: str | None = None
    time_limit: str | None = None

    Schema: ClassVar[type[Schema]] = Schema


@dataclass(frozen=True)
class BenchmarkConfig:
    """Benchmark configuration."""

    type: str = "manual"
    isl: int | None = None
    osl: int | None = None
    concurrencies: list[int] | str | None = None
    req_rate: str | int | None = "inf"
    sweep: Annotated[SweepConfig, SweepConfigField()] | None = None
    # Accuracy benchmark fields
    num_examples: int | None = None
    max_tokens: int | None = None
    repeat: int | None = None
    num_threads: int | None = None
    max_context_length: int | None = None
    categories: list[str] | None = None
    num_shots: int | None = None  # GSM8K few-shot examples
    aime_dataset: str | None = None  # NeMo Skills AIME dataset: aime24, aime25, or aime26
    temperature: float | None = None
    top_p: float | None = None
    top_k: int | None = None
    # Router benchmark fields
    num_requests: int | None = None
    concurrency: int | None = None
    prefix_ratios: list[float] | str | None = None
    # Mooncake router benchmark fields (uses aiperf with mooncake_trace)
    mooncake_workload: str | None = None  # "mooncake", "conversation", "synthetic", "toolagent"
    ttft_threshold_ms: int | None = None  # Goodput TTFT threshold in ms (default: 2000)
    itl_threshold_ms: int | None = None  # Goodput ITL threshold in ms (default: 25)
    random_range_ratio: float | None = None  # Random input/output length range ratio (default: 0.8)
    num_prompts_mult: int | None = None  # Multiplier for num_prompts = concurrency * mult (default: 10)
    num_warmup_mult: int | None = None  # Multiplier for warmup prompts = concurrency * mult (default: 2)
    # Custom dataset fields (sa-bench)
    dataset_name: str | None = None  # "random" (default) or "custom"
    dataset_path: str | None = None  # Container path to dataset file (mount via extra_mount)
    # Trace replay benchmark fields (uses aiperf with mooncake_trace dataset type)
    trace_file: str | None = None  # Path to trace JSONL file (container path, e.g., /traces/dataset.jsonl)
    custom_tokenizer: str | None = None  # Custom tokenizer class (e.g., "module.path.ClassName")
    use_chat_template: bool = True  # Pass --use-chat-template to benchmark (default: true)
    # Custom benchmark hook.
    # ``command`` is passed to ``bash -lc`` verbatim; srtctl does NOT
    # substitute placeholders like ``{nginx_url}`` or ``{slurm_job_id}``.
    # Render any parameters when generating the recipe. See
    # srtctl.benchmarks.custom.CustomBenchmarkRunner for details.
    command: str | None = None
    container_image: str | None = None
    env: dict[str, str] = field(default_factory=dict)
    # aiperf pip install spec (e.g., "aiperf>=0.7.0", "aiperf @ git+https://...@commit")
    # If set, runs pip install <spec> before benchmarking. Upgrades if already installed.
    aiperf_package: str | None = None
    # Extra aiperf CLI flags passed through to bench.sh (e.g., benchmark-duration: 600, workers-max: 200)
    aiperf_args: dict[str, Any] = field(default_factory=dict)

    def get_concurrency_list(self) -> list[int]:
        if self.concurrencies is None:
            return []
        if isinstance(self.concurrencies, str):
            return [int(x) for x in self.concurrencies.split("x")]
        return list(self.concurrencies)

    Schema: ClassVar[builtins.type[Schema]] = Schema


@dataclass(frozen=True)
class ProfilingPhaseConfig:
    """Profiling config for a single phase (prefill/decode/aggregated)."""

    start_step: int | None = None  # Step to start profiling
    stop_step: int | None = None  # Step to stop profiling

    Schema: ClassVar[builtins.type[Schema]] = Schema


@dataclass(frozen=True)
class ProfilingConfig:
    """Profiling configuration.

    Supports two profiling modes:
    - nsys: NVIDIA Nsight Systems profiling (wraps command with nsys profile)
    - torch: PyTorch profiler (uses SGLANG_TORCH_PROFILER_DIR)

    Per-phase start_step/stop_step are specified in the prefill/decode/aggregated sections.
    """

    type: str = "none"  # "none", "nsys", "nsys-time", or "torch"

    # Extra arguments passed to nsys profile (appended before `-o`; see get_nsys_prefix)
    extra_nsys_args: list[str] | None = None

    # Phase-specific profiling step configs (not used for nsys-time)
    prefill: ProfilingPhaseConfig | None = None
    decode: ProfilingPhaseConfig | None = None
    aggregated: ProfilingPhaseConfig | None = None

    # nsys-time fields: time-based capture window, same on all workers
    delay_secs: int | None = None  # nsys --delay: seconds from worker launch before capture starts
    duration_secs: int | None = None  # nsys --duration: seconds to capture after delay
    benchmark_duration_secs: int = 300  # total traffic generation duration (must cover delay + duration)

    @property
    def enabled(self) -> bool:
        """Check if profiling is enabled."""
        return self.type != "none"

    @property
    def is_nsys(self) -> bool:
        """Check if using NVIDIA Nsight Systems profiling (includes nsys-time)."""
        return self.type in ("nsys", "nsys-time")

    @property
    def is_nsys_time(self) -> bool:
        """Check if using time-based nsys capture (--delay/--duration instead of cudaProfilerApi)."""
        return self.type == "nsys-time"

    @property
    def is_torch(self) -> bool:
        """Check if using PyTorch profiler."""
        return self.type == "torch"

    def _get_phase_config(self, mode: str) -> ProfilingPhaseConfig | None:
        """Get the phase config for the given mode."""
        if mode == "prefill":
            return self.prefill
        elif mode == "decode":
            return self.decode
        elif mode in ("agg", "aggregated"):
            return self.aggregated
        return None

    def get_env_vars(self, mode: str, profile_dir: str) -> dict[str, str]:
        """Get profiling-specific environment variables.

        Args:
            mode: Worker mode (prefill/decode/agg)
            profile_dir: Base directory for profiling output

        Returns:
            Dictionary of environment variables
        """
        if not self.enabled:
            return {}

        env = {"PROFILING_MODE": mode, "PROFILE_TYPE": self.type}

        # Phase-specific start/stop steps
        phase_config = self._get_phase_config(mode)
        if phase_config:
            phase_key = mode.upper() if mode != "agg" else "AGG"
            if phase_config.start_step is not None:
                env[f"PROFILE_{phase_key}_START_STEP"] = str(phase_config.start_step)
            if phase_config.stop_step is not None:
                env[f"PROFILE_{phase_key}_STOP_STEP"] = str(phase_config.stop_step)

        if self.is_torch:
            env["SGLANG_TORCH_PROFILER_DIR"] = f"{profile_dir}/{mode}"

        if self.is_nsys_time:
            env["PROFILE_BENCHMARK_DURATION_SECS"] = str(self.benchmark_duration_secs)
        elif (
            self.is_nsys and phase_config and phase_config.start_step is not None and phase_config.stop_step is not None
        ):
            # TRTLLM iteration-based nsys: PyExecutor triggers cudaProfilerStart/Stop at these boundaries.
            # Harmless on SGLang workers (unknown env vars are ignored).
            env["TLLM_PROFILE_START_STOP"] = f"{phase_config.start_step}-{phase_config.stop_step}"
            env["TLLM_LLMAPI_ENABLE_NVTX"] = "1"

        return env

    def _get_nsys_prefix_trtllm(self, output_file: str) -> list[str]:
        """Get nsys command prefix for TRTLLM workers.

        Supports both iteration-based (cudaProfilerApi trigger via TLLM_PROFILE_START_STOP)
        and time-based (--delay/--duration) capture modes.
        """
        if self.is_nsys_time:
            cmd = [
                "nsys",
                "profile",
                "-t",
                "cuda,nvtx,ucx",
                "--sample=none",
                "--cuda-graph-trace=node",
            ]
            if self.delay_secs is not None:
                cmd += ["--delay", str(self.delay_secs)]
            if self.duration_secs is not None:
                cmd += ["--duration", str(self.duration_secs)]
        else:
            # Iteration-based: TLLM_PROFILE_START_STOP env var triggers cudaProfilerStart/Stop
            cmd = [
                "nsys",
                "profile",
                "-t",
                "cuda,nvtx,ucx",
                "--sample=none",
                "--cuda-graph-trace=node",
                "-c",
                "cudaProfilerApi",
                "--capture-range-end",
                "stop",
            ]

        if self.extra_nsys_args:
            cmd.extend(self.extra_nsys_args)

        cmd += [
            "--kill",
            "none",
            "--wait",
            "all",
            "--force-overwrite",
            "true",
            "-o",
            output_file,
        ]
        return cmd

    def get_nsys_prefix(
        self, output_file: str, *, frontend_type: str | None = None, backend_type: str | None = None
    ) -> list[str]:
        """Get nsys profiling command prefix.

        Args:
            output_file: Path for nsys output file (without extension)
            frontend_type: Frontend type (e.g., "dynamo", "sglang"). When set to "dynamo"
                with a non-trtllm backend, adds --trace-fork-before-exec=true.
            backend_type: Backend type (e.g., "trtllm", "sglang"). When set to "trtllm",
                uses TRTLLM-specific nsys flags (ucx traces, --kill none, --wait all).

        Returns:
            Command prefix list for nsys profiling
        """
        if not self.is_nsys:
            return []

        if backend_type == "trtllm":
            return self._get_nsys_prefix_trtllm(output_file)

        # SGLang / default path — keep existing behavior
        cmd = [
            "nsys",
            "profile",
            "-t",
            "cuda,nvtx",
            "--cuda-graph-trace=node",
            "-c",
            "cudaProfilerApi",
            "--capture-range-end",
            "stop",
            "--force-overwrite",
            "true",
        ]

        if self.extra_nsys_args:
            cmd.extend(self.extra_nsys_args)

        cmd.extend(["-o", output_file])

        if frontend_type == "dynamo":
            cmd.insert(-2, "--trace-fork-before-exec=true")

        return cmd

    Schema: ClassVar[builtins.type[Schema]] = Schema


@dataclass(frozen=True)
class ObservabilityConfig:
    """Observability configuration for OTEL tracing.

    When enable_otel is True, OTEL environment variables (DYN_LOGGING_JSONL,
    OTEL_EXPORT_ENABLED, OTEL_EXPORTER_OTLP_TRACES_ENDPOINT, OTEL_SERVICE_NAME)
    are automatically injected into all workers and frontends.

    OTEL_SERVICE_NAME defaults to "dynamo-{component}" (e.g. dynamo-prefill,
    dynamo-decode, dynamo-frontend) and can be overridden per-component via
    prefill_environment, decode_environment, or frontend.env.

    Attributes:
        enable_otel: If True, inject OTEL environment variables into all workers
            and frontends. Requires otel_endpoint to be set. Default: False.
        otel_endpoint: OTEL collector endpoint (e.g. "http://10.0.0.1:4317").
            Required when enable_otel is True.
    """

    enable_otel: bool = False
    otel_endpoint: str | None = None

    Schema: ClassVar[type[Schema]] = Schema


@dataclass(frozen=True)
class TelemetryExporterConfig:
    """Configuration for telemetry exporters deployed on worker nodes."""

    container_image: str
    port: int
    command: str | None = None

    Schema: ClassVar[type[Schema]] = Schema


@dataclass(frozen=True)
class TelemetryConfig:
    """Telemetry configuration for benchmark jobs.

    The default provider bundles a scraper with dcgm_exporter and node_exporter.
    Other providers can reuse the same top-level contract later.
    """

    enabled: bool = False
    provider: TelemetryProvider = TelemetryProvider.SCRAPER
    container_image: str | None = None
    binary_path: str = "/usr/local/bin/telemetry-scraper"
    default_frequency: float = 5.0
    sync_interval_secs: int = 120
    compaction_threads: int = 4
    storage_subdir: str = "telemetry"
    extra_metadata: dict[str, str] = field(default_factory=dict)
    dcgm_exporter: TelemetryExporterConfig | None = None
    node_exporter: TelemetryExporterConfig | None = None

    Schema: ClassVar[type[Schema]] = Schema


def build_otel_env(observability: ObservabilityConfig, component: str) -> dict[str, str]:
    """Build OTEL environment variables for a component.

    Returns an empty dict if OTEL is disabled. Otherwise returns env vars
    with OTEL_SERVICE_NAME set to "dynamo-{component}".
    """
    if not observability.enable_otel or not observability.otel_endpoint:
        return {}
    return {
        "DYN_LOGGING_JSONL": "1",
        "OTEL_EXPORT_ENABLED": "1",
        "OTEL_EXPORTER_OTLP_TRACES_ENDPOINT": observability.otel_endpoint,
        "OTEL_SERVICE_NAME": f"dynamo-{component}",
    }


@dataclass
class DynamoConfig:
    """Dynamo installation configuration.

    Only one of version, hash, or top_of_tree should be specified.
    Defaults to version="0.8.0" (pip install).

    Options:
        install: Whether to install dynamo at all (default: True). Set to False
                 if your container already has dynamo pre-installed.
        version: Install specific version from PyPI (e.g., "0.8.0")
        hash: Clone repo and checkout specific commit hash
        top_of_tree: Clone repo at HEAD (latest)

    If top_of_tree or hash is set, version is automatically cleared.
    """

    install: bool = True
    version: str | None = "0.8.0"
    hash: str | None = None
    top_of_tree: bool = False

    def __post_init__(self) -> None:
        # Auto-clear version if hash or top_of_tree is set
        if self.hash is not None or self.top_of_tree:
            object.__setattr__(self, "version", None)

        # Validate only one source option is set
        if self.hash is not None and self.top_of_tree:
            raise ValueError("Cannot specify both hash and top_of_tree")

    @property
    def needs_source_install(self) -> bool:
        """Whether this config requires a source install (git clone + maturin)."""
        return self.hash is not None or self.top_of_tree

    def get_install_commands(self) -> str:
        """Get the bash commands to install dynamo."""
        if self.version is not None:
            return (
                f"echo 'Installing dynamo {self.version}...' && "
                f"pip install --break-system-packages --quiet --extra-index-url https://pypi.nvidia.com ai-dynamo-runtime=={self.version} ai-dynamo=={self.version} && "
                f"echo 'Dynamo {self.version} installed'"
            )

        # Source install (hash or top-of-tree)
        git_ref = self.hash if self.hash else "HEAD"
        checkout_cmd = f"git checkout {self.hash}" if self.hash else ""

        # Original SGLang container path
        sglang = (
            # protobuf-compiler is required by modelexpress-common's build.rs (prost-build).
            # Some SGLang images ship without /usr/bin/protoc; install it unconditionally.
            "apt-get update -qq && apt-get install -y -qq libclang-dev curl protobuf-compiler > /dev/null 2>&1 && "
            "if ! command -v cargo &>/dev/null; then curl --proto '=https' --tlsv1.2 -sSf https://sh.rustup.rs | sh -s -- -y --default-toolchain stable -q && source $HOME/.cargo/env; fi && "
            # Force-reinstall maturin: some images ship the python module in dist-packages
            # without the console-script entry point, so `command -v maturin` fails AND a
            # plain `pip install maturin` reports "already satisfied" and skips the fix.
            "pip install --break-system-packages --force-reinstall --quiet maturin && "
            "cd /sgl-workspace/ && "
            "git clone https://github.com/ai-dynamo/dynamo.git && "
            "cd dynamo && "
            f"{checkout_cmd + ' && ' if checkout_cmd else ''}"
            "cd lib/bindings/python/ && "
            'export RUSTFLAGS="${RUSTFLAGS:-} -C target-cpu=native --cfg tokio_unstable" && '
            "maturin build -o /tmp && "
            "pip install /tmp/ai_dynamo_runtime*.whl && "
            "cd /sgl-workspace/dynamo/ && "
            "pip install -e . && "
            "cd /sgl-workspace/sglang/ && "
            f"echo 'Dynamo installed from source ({git_ref})'"
        )

        # Portable path for non-SGLang containers (vLLM, etc.)
        portable = (
            "if ! command -v cargo &> /dev/null || ! command -v maturin &> /dev/null; then "
            "apt-get update -qq && apt-get install -y -qq git curl libclang-dev protobuf-compiler > /dev/null 2>&1 && "
            "if ! command -v cargo &> /dev/null; then "
            "curl --proto '=https' --tlsv1.2 -sSf https://sh.rustup.rs | sh -s -- -y && source $HOME/.cargo/env; fi && "
            "if ! command -v maturin &> /dev/null; then "
            "pip install --break-system-packages maturin; fi; fi && "
            "ORIG_DIR=$(pwd) && rm -rf /tmp/dynamo_build && mkdir -p /tmp/dynamo_build && cd /tmp/dynamo_build && "
            "git clone https://github.com/ai-dynamo/dynamo.git && "
            "cd dynamo && "
            f"{checkout_cmd + ' && ' if checkout_cmd else ''}"
            "cd lib/bindings/python/ && "
            'export RUSTFLAGS="${RUSTFLAGS:-} -C target-cpu=native --cfg tokio_unstable" && '
            "rm -f /tmp/ai_dynamo_runtime*.whl && "
            "maturin build -o /tmp && "
            "pip install --break-system-packages /tmp/ai_dynamo_runtime*.whl --force-reinstall && "
            "cd /tmp/dynamo_build/dynamo/ && "
            "pip install --break-system-packages -e . && "
            "cd $ORIG_DIR && "
            f"echo 'Dynamo installed from source ({git_ref})'"
        )

        return (
            f"echo 'Installing dynamo from source ({git_ref})...' && "
            f"if [ -d /sgl-workspace ]; then {sglang}; else {portable}; fi"
        )

    Schema: ClassVar[type[Schema]] = Schema


@dataclass(frozen=True)
class FrontendConfig:
    """Frontend/router configuration.

    Attributes:
        type: Frontend type - "dynamo" (default) or "sglang"
        enable_multiple_frontends: Scale with nginx + multiple routers.
            When ``True`` (default), srtctl stands up nginx and fans out
            to ``num_additional_frontends + 1`` router replicas. When
            ``False``, there is NO nginx proxy — the benchmark must
            target the single master router (or a worker) directly at
            ``http://localhost:<port>``. ``benchmark.command`` has no
            placeholder substitution, so write the URL out literally.
        num_additional_frontends: Additional routers beyond master (default: 9)
        nginx_container: Custom nginx container image (default: nginx:1.27.4)
        args: CLI arguments passed to the frontend/router process
        env: Environment variables for frontend processes
    """

    type: str = "dynamo"
    enable_multiple_frontends: bool = True
    num_additional_frontends: int = 9
    nginx_container: str = "nginx:1.27.4"
    args: dict[str, Any] | None = None
    env: dict[str, str] | None = None

    Schema: ClassVar[builtins.type[Schema]] = Schema


@dataclass(frozen=True)
class OutputConfig:
    """Output configuration with formattable paths."""

    log_dir: Annotated[FormattablePath, FormattablePathField()] = field(
        default_factory=lambda: FormattablePath(template="./outputs/{job_id}/logs")
    )

    Schema: ClassVar[type[Schema]] = Schema


@dataclass(frozen=True)
class HealthCheckConfig:
    """Health check configuration."""

    max_attempts: int = 180  # 30 minutes default (large models take time to load)
    interval_seconds: int = 10

    Schema: ClassVar[type[Schema]] = Schema


@dataclass(frozen=True)
class InfraConfig:
    """Infrastructure configuration for etcd/nats placement.

    Attributes:
        etcd_nats_dedicated_node: If True, run etcd and nats on a dedicated node
            instead of the head node. This reserves the first node exclusively
            for infrastructure services. Default: False.
        nats_max_payload_mb: Maximum NATS message payload in MB. Default: None (uses
            NATS default of 1MB). Set to 24+ for disaggregated serving with long ISL
            (e.g. 65K+ tokens where prompt data exceeds 1MB in NATS messages).
    """

    etcd_nats_dedicated_node: bool = False
    nats_max_payload_mb: int | None = None

    Schema: ClassVar[type[Schema]] = Schema


# ============================================================================
# Main Configuration Dataclass
# ============================================================================


@dataclass(frozen=True)
class SrtConfig:
    """Complete srtctl job configuration (frozen, immutable).

    This is the main configuration type returned by load_config().

    The backend field supports polymorphic deserialization:
    - type: sglang -> SGLangProtocol
    """

    name: str
    model: ModelConfig
    resources: ResourceConfig

    slurm: SlurmConfig = field(default_factory=SlurmConfig)
    backend: Annotated[BackendConfig, BackendConfigField()] = field(default_factory=SGLangProtocol)
    frontend: FrontendConfig = field(default_factory=FrontendConfig)
    dynamo: DynamoConfig = field(default_factory=DynamoConfig)
    benchmark: BenchmarkConfig = field(default_factory=BenchmarkConfig)
    profiling: ProfilingConfig = field(default_factory=ProfilingConfig)
    output: OutputConfig = field(default_factory=OutputConfig)
    health_check: HealthCheckConfig = field(default_factory=HealthCheckConfig)
    infra: InfraConfig = field(default_factory=InfraConfig)
    observability: ObservabilityConfig = field(default_factory=ObservabilityConfig)
    telemetry: TelemetryConfig = field(default_factory=TelemetryConfig)

    environment: dict[str, str] = field(default_factory=dict)
    container_mounts: dict[
        Annotated[FormattablePath, FormattablePathField()],
        Annotated[FormattablePath, FormattablePathField()],
    ] = field(default_factory=dict)
    extra_mount: tuple[str, ...] | None = None
    srun_options: dict[str, str] = field(default_factory=dict)
    sbatch_directives: dict[str, str] = field(default_factory=dict)
    enable_config_dump: bool = True

    # Custom setup script (runs before dynamo install and worker startup)
    # e.g. "custom-setup.sh" -> runs /configs/custom-setup.sh
    setup_script: str | None = None

    # Virtual identity — declares what *should* be running (verified against fingerprint)
    identity: IdentityConfig = field(default_factory=IdentityConfig)

    # Reporting configuration (status API, future: logs to S3, etc.)
    reporting: ReportingConfig | None = None

    Schema: ClassVar[type[Schema]] = Schema

    def __post_init__(self):
        """Validate configuration after initialization."""
        self._validate_profiling()
        self._validate_telemetry()

    def _validate_profiling(self):
        """Validate profiling configuration matches serving mode."""
        prof = self.profiling
        if not prof.enabled:
            return

        backend_type = self.backend.type

        # torch profiling is SGLang-only (uses SGLANG_TORCH_PROFILER_DIR)
        if prof.is_torch and backend_type == "trtllm":
            raise ValidationError("torch profiling is not supported for the trtllm backend; use nsys instead")

        # nsys-time is TRTLLM-only (time-based capture via nsys --delay/--duration)
        if prof.is_nsys_time and backend_type != "trtllm":
            raise ValidationError("nsys-time profiling is only supported for the trtllm backend")

        # nsys-time uses top-level delay/duration — no per-phase step configs needed
        if prof.is_nsys_time:
            if prof.delay_secs is None or prof.duration_secs is None:
                raise ValidationError(
                    "profiling.delay_secs and profiling.duration_secs are required for nsys-time mode"
                )
            return

        r = self.resources
        is_disaggregated = r.is_disaggregated
        has_prefill_prof = prof.prefill is not None
        has_decode_prof = prof.decode is not None
        has_agg_prof = prof.aggregated is not None

        # Validate phase configs match serving mode
        if is_disaggregated:
            if has_agg_prof:
                raise ValidationError(
                    "Disaggregated mode only supports profiling.prefill/decode; profiling.aggregated is not allowed."
                )
            if not has_prefill_prof or not has_decode_prof:
                raise ValidationError(
                    "Disaggregated mode requires both profiling.prefill and profiling.decode "
                    "to be set when profiling is enabled."
                )
            if (r.prefill_workers or 0) <= 0 or (r.decode_workers or 0) <= 0:
                raise ValidationError("Disaggregated mode requires prefill_workers and decode_workers to be > 0.")
        else:
            if has_prefill_prof or has_decode_prof:
                raise ValidationError(
                    "Aggregated mode only supports profiling.aggregated; profiling.prefill/decode are not allowed."
                )
            if not has_agg_prof:
                raise ValidationError(
                    "Aggregated mode requires profiling.aggregated to be set when profiling is enabled."
                )
            if (r.agg_workers or 0) <= 0:
                raise ValidationError("Aggregated mode requires agg_workers to be > 0.")

    def _validate_telemetry(self):
        """Validate telemetry configuration."""
        telemetry = self.telemetry
        if not telemetry.enabled:
            return

        if telemetry.provider != TelemetryProvider.SCRAPER:
            raise ValidationError(f"Unsupported telemetry provider: {telemetry.provider}")

        if not telemetry.container_image:
            raise ValidationError("telemetry.container_image is required when telemetry is enabled")
        if telemetry.dcgm_exporter is None:
            raise ValidationError("telemetry.dcgm_exporter is required when telemetry is enabled")
        if telemetry.node_exporter is None:
            raise ValidationError("telemetry.node_exporter is required when telemetry is enabled")
        if telemetry.default_frequency <= 0:
            raise ValidationError("telemetry.default_frequency must be positive")
        if telemetry.sync_interval_secs < 0:
            raise ValidationError("telemetry.sync_interval_secs must be >= 0")
        if telemetry.compaction_threads < 0:
            raise ValidationError("telemetry.compaction_threads must be >= 0")

    @classmethod
    def from_yaml(cls, yaml_path: Path) -> "SrtConfig":
        with open(yaml_path) as f:
            data = yaml.safe_load(f)
        schema = cls.Schema()
        return schema.load(data)

    @property
    def served_model_name(self) -> str:
        """Get the served model name from backend config or model path."""
        default = Path(self.model.path).name
        return self.backend.get_served_model_name(default)

    @property
    def backend_type(self) -> str:
        """Get the backend type string."""
        return self.backend.type