diff --git a/packages/data-designer-engine/src/data_designer/engine/models/telemetry.py b/packages/data-designer-engine/src/data_designer/engine/models/telemetry.py index 41f0b9c18..b2b01384b 100644 --- a/packages/data-designer-engine/src/data_designer/engine/models/telemetry.py +++ b/packages/data-designer-engine/src/data_designer/engine/models/telemetry.py @@ -51,6 +51,7 @@ class NemoSourceEnum(str, Enum): class DeploymentTypeEnum(str, Enum): LIBRARY = "library" API = "api" + NVIDIA_INTERNAL = "nvidia-internal" UNDEFINED = "undefined" @@ -72,7 +73,7 @@ class TaskStatusEnum(str, Enum): class TelemetryEvent(BaseModel): _event_name: ClassVar[str] # Subclasses must define this - _schema_version: ClassVar[str] = "1.3" + _schema_version: ClassVar[str] = "1.9" def __init_subclass__(cls, **kwargs: Any) -> None: super().__init_subclass__(**kwargs) diff --git a/packages/data-designer-engine/tests/engine/models/test_telemetry.py b/packages/data-designer-engine/tests/engine/models/test_telemetry.py new file mode 100644 index 000000000..f78739097 --- /dev/null +++ b/packages/data-designer-engine/tests/engine/models/test_telemetry.py @@ -0,0 +1,32 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +from datetime import datetime, timezone + +from data_designer.engine.models.telemetry import ( + DeploymentTypeEnum, + InferenceEvent, + NemoSourceEnum, + QueuedEvent, + TaskStatusEnum, + build_payload, +) + + +def test_nvidia_internal_deployment_type_uses_schema_version_1_9() -> None: + event = InferenceEvent( + nemo_source=NemoSourceEnum.DATADESIGNER, + task="batch", + task_status=TaskStatusEnum.SUCCESS, + deployment_type=DeploymentTypeEnum("nvidia-internal"), + model="test-model", + ) + payload = build_payload( + [QueuedEvent(event=event, timestamp=datetime(2026, 7, 2, tzinfo=timezone.utc))], + source_client_version="test", + ) + + assert payload["eventSchemaVer"] == "1.9" + assert payload["events"][0]["parameters"]["deploymentType"] == "nvidia-internal"