schema.py

# SPDX-FileCopyrightText: 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.
"""Pydantic models for GXT Event Protocol v1.6 telemetry payloads.

These models define the wire format for TRT-LLM usage telemetry sent to
the NvTelemetry/GXT endpoint. The envelope follows the GXT Event Protocol v1.6
specification; the event parameters are TRT-LLM-specific.

Reference:
- GXT API Design: GX Telemetry API Design.md
- DataDesigner reference: DataDesigner telemetry.py
"""

import platform
from datetime import datetime, timezone
from typing import Any, Dict, List, Literal, Optional, Union

from pydantic import BaseModel, Field

# ---------------------------------------------------------------------------
# Constants
# ---------------------------------------------------------------------------

_UINT32_MAX = 4_294_967_295  # 2**32 - 1; NvTelemetry "PositiveInt"
_SHORT_STR = 128  # NvTelemetry "ShortString" maxLength
_LONG_STR = 256  # NvTelemetry "LongString" maxLength

CLIENT_ID = "616561816355034"
EVENT_PROTOCOL = "1.6"
EVENT_SCHEMA_VER = "0.7"
EVENT_SYS_VER = "trtllm-telemetry/1.0"
CLIENT_TYPE = "Native"
CLIENT_VARIANT = "Release"
CPU_ARCHITECTURE = platform.uname().machine


# ---------------------------------------------------------------------------
# TRT-LLM Event Parameters (inner payloads)
# ---------------------------------------------------------------------------


class _LlmCounterSnapshot(BaseModel):
    """Process-local aggregate LLM lifecycle counters.

    - Initialization attempts: entries into LLM construction.
    - Instances created: LLM objects initialized successfully.
    - Active instances: initialized objects not yet shut down.
    - Maximum concurrent instances: highest active count observed.
    - Initialization failures: construction attempts that raised an exception.
    """

    llm_initialization_attempts: int = Field(
        default=0,
        ge=0,
        le=_UINT32_MAX,
        alias="llmInitializationAttempts",
        description="Number of LLM construction attempts in this process.",
    )
    llm_instances_created: int = Field(
        default=0,
        ge=0,
        le=_UINT32_MAX,
        alias="llmInstancesCreated",
        description="Number of LLM instances successfully created in this process.",
    )
    active_llm_instances: int = Field(
        default=0,
        ge=0,
        le=_UINT32_MAX,
        alias="activeLlmInstances",
        description="Number of initialized LLM instances not yet shut down.",
    )
    max_concurrent_llm_instances: int = Field(
        default=0,
        ge=0,
        le=_UINT32_MAX,
        alias="maxConcurrentLlmInstances",
        description="Highest number of concurrently active LLM instances observed.",
    )
    llm_initialization_failures: int = Field(
        default=0,
        ge=0,
        le=_UINT32_MAX,
        alias="llmInitializationFailures",
        description="Number of LLM construction attempts that raised an exception.",
    )

    model_config = {"populate_by_name": True}


class TrtllmInitialReport(_LlmCounterSnapshot):
    """TRT-LLM initial report event parameters.

    Sent once at startup with full environment and configuration details.
    All fields are required by the SMS schema (GXT convention: every declared
    property must be in ``required``). Fields use sentinel defaults (empty
    string for strings, 0 for ints) when the actual value is unavailable.

    Field constraints match the SMS JSON schema type definitions:
    - ShortString: maxLength=128
    - LongString: maxLength=256
    - PositiveInt: ge=0, le=_UINT32_MAX (2**32 - 1)
    """

    # TRT-LLM version (ShortString)
    trtllm_version: str = Field(default="", max_length=_SHORT_STR, alias="trtllmVersion")

    # System info
    platform_info: str = Field(default="", max_length=_LONG_STR, alias="platform")  # LongString
    python_version: str = Field(
        default="", max_length=_SHORT_STR, alias="pythonVersion"
    )  # ShortString
    cpu_architecture: str = Field(
        default="", max_length=_SHORT_STR, alias="cpuArchitecture"
    )  # ShortString
    cpu_count: int = Field(default=0, ge=0, le=_UINT32_MAX, alias="cpuCount")  # PositiveInt

    # GPU info
    gpu_count: int = Field(default=0, ge=0, le=_UINT32_MAX, alias="gpuCount")  # PositiveInt
    gpu_name: str = Field(default="", max_length=_LONG_STR, alias="gpuName")  # LongString
    gpu_memory_mb: int = Field(default=0, ge=0, le=_UINT32_MAX, alias="gpuMemoryMB")  # PositiveInt
    cuda_version: str = Field(default="", max_length=_SHORT_STR, alias="cudaVersion")  # ShortString

    # Model info (architecture class name only -- no raw config) (LongString)
    architecture_class_name: str = Field(
        default="", max_length=_LONG_STR, alias="architectureClassName"
    )
    architecture_class_hash: str = Field(
        default="",
        max_length=_LONG_STR,
        alias="architectureClassHash",
        description="Stable hash of the model architecture class definition.",
    )

    # TRT-LLM config
    backend: str = Field(default="", max_length=_SHORT_STR, alias="backend")  # ShortString
    tensor_parallel_size: int = Field(default=1, ge=0, le=_UINT32_MAX, alias="tensorParallelSize")
    pipeline_parallel_size: int = Field(
        default=1, ge=0, le=_UINT32_MAX, alias="pipelineParallelSize"
    )
    context_parallel_size: int = Field(default=1, ge=0, le=_UINT32_MAX, alias="contextParallelSize")
    moe_expert_parallel_size: int = Field(
        default=0, ge=0, le=_UINT32_MAX, alias="moeExpertParallelSize"
    )
    moe_tensor_parallel_size: int = Field(
        default=0, ge=0, le=_UINT32_MAX, alias="moeTensorParallelSize"
    )
    dtype: str = Field(default="", max_length=_SHORT_STR, alias="dtype")  # ShortString
    quantization_algo: str = Field(
        default="", max_length=_SHORT_STR, alias="quantizationAlgo"
    )  # ShortString
    kv_cache_dtype: str = Field(
        default="", max_length=_SHORT_STR, alias="kvCacheDtype"
    )  # ShortString

    # Ingress point (how TRT-LLM was invoked) (ShortString)
    ingress_point: str = Field(
        default="",
        max_length=_SHORT_STR,
        alias="ingressPoint",
        description="TRT-LLM API or CLI entry point that created this session.",
    )

    # Disaggregated serving metadata (ShortString)
    disagg_role: str = Field(
        default="",
        max_length=_SHORT_STR,
        alias="disaggRole",
        description="Role of this process in a disaggregated deployment.",
    )
    deployment_id: str = Field(
        default="",
        max_length=_SHORT_STR,
        alias="deploymentId",
        description="Identifier used to correlate a disaggregated deployment.",
    )

    # Legacy feature summary and sanitized opt-in LLM API config capture.
    features_json: str = Field(default="{}", alias="featuresJson")
    llm_api_config_json: str = Field(default="{}", alias="llmApiConfigJson")
    llm_api_config_meta_json: str = Field(default="{}", alias="llmApiConfigMetaJson")

    model_config = {"populate_by_name": True}


class TrtllmHeartbeat(_LlmCounterSnapshot):
    """TRT-LLM heartbeat event parameters.

    Sent periodically to signal the session is still alive.
    Contains a monotonically increasing sequence counter, process correlation
    fields, and the latest aggregate LLM lifecycle counter snapshot.
    """

    seq: int = Field(
        ...,
        ge=0,
        le=_UINT32_MAX,
        alias="seq",
        description="Zero-based heartbeat sequence number for this session.",
    )
    ingress_point: str = Field(
        default="",
        max_length=_SHORT_STR,
        alias="ingressPoint",
        description="TRT-LLM API or CLI entry point that created this session.",
    )
    disagg_role: str = Field(
        default="",
        max_length=_SHORT_STR,
        alias="disaggRole",
        description="Role of this process in a disaggregated deployment.",
    )
    deployment_id: str = Field(
        default="",
        max_length=_SHORT_STR,
        alias="deploymentId",
        description="Identifier used to correlate a disaggregated deployment.",
    )

    model_config = {"populate_by_name": True}


TerminationKind = Literal[
    "clean",
    "exception",
    "signal",
    "worker_failure",
    "timeout",
    "unknown",
]
LifecyclePhase = Literal[
    "cli_parsing",
    "config_validation",
    "model_initialization",
    "serving",
    "unknown",
]
TerminationComponent = Literal[
    "llm",
    "server",
    "engine_worker",
    "disagg_worker",
    "unknown",
]
ReportingSource = Literal["self", "supervisor", "executor_proxy"]


class TrtllmExitReport(_LlmCounterSnapshot):
    """TRT-LLM terminal event parameters.

    Sent at most once for a telemetry session when TRT-LLM or a surviving
    observer can classify the session outcome. Unknown values use explicit
    sentinels instead of inferred causes.
    """

    exit_code_known: bool = Field(
        ...,
        alias="exitCodeKnown",
        description="Whether exitCode is an observed process exit code.",
    )
    exit_code: int = Field(
        ...,
        ge=0,
        le=_UINT32_MAX,
        alias="exitCode",
        description="Observed process exit code, or zero when unavailable.",
    )
    signal_number: int = Field(
        ...,
        ge=0,
        le=_UINT32_MAX,
        alias="signalNumber",
        description="Observed terminating signal number, or zero when unavailable.",
    )
    termination_kind: TerminationKind = Field(
        ...,
        alias="terminationKind",
        description="Allowlisted classification of how the session terminated.",
    )
    lifecycle_phase: LifecyclePhase = Field(
        ...,
        alias="lifecyclePhase",
        description="Last known lifecycle phase when the session terminated.",
    )
    component: TerminationComponent = Field(
        ...,
        alias="component",
        description="TRT-LLM component associated with the terminal outcome.",
    )
    reporting_source: ReportingSource = Field(
        ...,
        alias="reportingSource",
        description="Observer that classified and reported the terminal outcome.",
    )
    ingress_point: str = Field(
        default="",
        max_length=_SHORT_STR,
        alias="ingressPoint",
        description="TRT-LLM API or CLI entry point that created this session.",
    )
    disagg_role: str = Field(
        default="",
        max_length=_SHORT_STR,
        alias="disaggRole",
        description="Role of this process in a disaggregated deployment.",
    )
    deployment_id: str = Field(
        default="",
        max_length=_SHORT_STR,
        alias="deploymentId",
        description="Identifier used to correlate a disaggregated deployment.",
    )

    model_config = {"populate_by_name": True}


# ---------------------------------------------------------------------------
# GXT Event Wrapper (single event in the events array)
# ---------------------------------------------------------------------------


class GxtEvent(BaseModel):
    """A single event entry in the GXT events array."""

    ts: str = Field(..., description="ISO 8601 timestamp")
    name: str = Field(default="trtllm_usage_event")
    parameters: Dict[str, Any] = Field(...)


# ---------------------------------------------------------------------------
# GXT Envelope (top-level payload)
# ---------------------------------------------------------------------------


class GxtPayload(BaseModel):
    """GXT Event Protocol v1.6 envelope.

    The GXT ingestion endpoint validates a fixed envelope schema and rejects
    payloads that omit any of the required top-level fields — even fields
    that are semantically irrelevant for this client.

    GXT was designed for consumer products (GeForce Experience, NVIDIA App)
    where device fingerprints, user IDs, and GDPR consent flags carry real
    values. TRT-LLM is a server-side SDK with no browser, device, or login,
    so those fields are hardcoded to sentinel values:

    - Privacy/identity fields → "undefined"
      Signals to downstream pipelines and auditors that these fields were
      *deliberately not collected*, not accidentally missing.

    - GDPR opt-in fields → "None" (the string, not Python None)
      In GXT's data-policy framework this means "no consent model applies".

    This pattern matches the DataDesigner reference implementation (another
    internal GXT client for server-side telemetry).
    """

    # Client identification
    client_id: str = Field(default=CLIENT_ID, alias="clientId")
    client_type: str = Field(default=CLIENT_TYPE, alias="clientType")
    client_variant: str = Field(default=CLIENT_VARIANT, alias="clientVariant")
    client_ver: str = Field(..., alias="clientVer")
    cpu_architecture: str = Field(default=CPU_ARCHITECTURE, alias="cpuArchitecture")

    # Protocol metadata
    event_protocol: str = Field(default=EVENT_PROTOCOL, alias="eventProtocol")
    event_schema_ver: str = Field(default=EVENT_SCHEMA_VER, alias="eventSchemaVer")
    event_sys_ver: str = Field(default=EVENT_SYS_VER, alias="eventSysVer")

    # Session
    sent_ts: str = Field(..., alias="sentTs")
    session_id: str = Field(..., alias="sessionId")

    # Required by GXT schema but unused by server-side SDK clients.
    # "undefined" = deliberately not collected (not accidentally missing).
    browser_type: str = Field(default="undefined", alias="browserType")
    device_id: str = Field(default="undefined", alias="deviceId")
    device_make: str = Field(default="undefined", alias="deviceMake")
    device_model: str = Field(default="undefined", alias="deviceModel")
    device_os: str = Field(default="undefined", alias="deviceOS")
    device_os_version: str = Field(default="undefined", alias="deviceOSVersion")
    device_type: str = Field(default="undefined", alias="deviceType")
    user_id: str = Field(default="undefined", alias="userId")
    external_user_id: str = Field(default="undefined", alias="externalUserId")
    idp_id: str = Field(default="undefined", alias="idpId")
    integration_id: str = Field(default="undefined", alias="integrationId")
    product_name: str = Field(default="undefined", alias="productName")
    product_version: str = Field(default="undefined", alias="productVersion")

    # Required by GXT schema but no consent model applies for opt-out
    # server-side telemetry. "None" (string) = no GDPR consent tracking.
    gdpr_beh_opt_in: str = Field(default="None", alias="gdprBehOptIn")
    gdpr_func_opt_in: str = Field(default="None", alias="gdprFuncOptIn")
    gdpr_tech_opt_in: str = Field(default="None", alias="gdprTechOptIn")
    device_gdpr_beh_opt_in: str = Field(default="None", alias="deviceGdprBehOptIn")
    device_gdpr_func_opt_in: str = Field(default="None", alias="deviceGdprFuncOptIn")
    device_gdpr_tech_opt_in: str = Field(default="None", alias="deviceGdprTechOptIn")

    # Events array
    events: List[GxtEvent] = Field(...)

    model_config = {"populate_by_name": True}


# ---------------------------------------------------------------------------
# Helper: build a ready-to-serialize payload
# ---------------------------------------------------------------------------


def get_iso_timestamp(dt: Optional[datetime] = None) -> str:
    """Return ISO 8601 timestamp with millisecond precision."""
    if dt is None:
        dt = datetime.now(timezone.utc)
    return dt.strftime("%Y-%m-%dT%H:%M:%S.") + f"{dt.microsecond // 1000:03d}Z"


def build_gxt_payload(
    event: Union[TrtllmInitialReport, TrtllmHeartbeat, TrtllmExitReport],
    *,
    session_id: str,
    trtllm_version: str,
) -> dict:
    """Build a complete GXT payload dict ready for json.dumps().

    Args:
        event: The TRT-LLM event to send.
        session_id: Ephemeral session UUID (hex string).
        trtllm_version: TRT-LLM package version string.

    Returns:
        dict suitable for json.dumps() and HTTP POST body.
    """
    now = get_iso_timestamp()

    if isinstance(event, TrtllmInitialReport):
        event_name = "trtllm_initial_report"
    elif isinstance(event, TrtllmHeartbeat):
        event_name = "trtllm_heartbeat"
    elif isinstance(event, TrtllmExitReport):
        event_name = "trtllm_exit_report"
    else:
        raise TypeError(f"Unknown event type: {type(event).__name__}")

    gxt_event = GxtEvent(
        ts=now,
        name=event_name,
        parameters=event.model_dump(by_alias=True),
    )

    payload = GxtPayload(
        client_ver=trtllm_version,
        sent_ts=now,
        session_id=session_id,
        events=[gxt_event],
    )

    return payload.model_dump(by_alias=True)