Merge branch 'main' into johnny/docs/plugins-new-structure-and-seed-type

johnnygreco · web-flow · commit c2f1dddf1fc2 · 2026-02-06T14:04:32.000-05:00
diff --git a/.github/workflows/health-checks.yml b/.github/workflows/health-checks.yml
@@ -0,0 +1,35 @@
+name: Health Checks
+
+on:
+  schedule:
+    # Every Monday at 08:00 UTC
+    - cron: "0 8 * * 1"
+  release:
+    types: [published]
+  workflow_dispatch:
+
+jobs:
+  health-checks:
+    name: Provider Health Checks
+    runs-on: ubuntu-latest
+
+    steps:
+      - name: Checkout code
+        uses: actions/checkout@v4
+
+      - name: Install uv
+        uses: astral-sh/setup-uv@v5
+        with:
+          version: "latest"
+          python-version: "3.11"
+          enable-cache: true
+
+      - name: Install dependencies
+        run: make install-dev
+
+      - name: Run health checks
+        env:
+          NVIDIA_API_KEY: ${{ secrets.TEST_NVIDIA_API_KEY }}
+          OPENAI_API_KEY: ${{ secrets.TEST_OPENAI_API_KEY }}
+          OPENROUTER_API_KEY: ${{ secrets.TEST_OPENROUTER_API_KEY }}
+        run: make health-checks
diff --git a/Makefile b/Makefile
@@ -48,6 +48,7 @@ help:
 	@echo "  test                      - Run all unit tests"
 	@echo "  coverage                  - Run tests with coverage report"
 	@echo "  test-e2e                  - Run e2e plugin tests"
+	@echo "  health-checks             - Run provider health checks"
 	@echo "  test-run-tutorials        - Run tutorial notebooks as e2e tests"
 	@echo "  test-run-recipes          - Run recipe scripts as e2e tests"
 	@echo "  test-run-all-examples     - Run all tutorials and recipes as e2e tests"
@@ -273,6 +274,10 @@ test-e2e:
 	@echo "🧪 Running e2e tests..."
 	uv run --no-cache --refresh --directory tests_e2e pytest -s
 
+health-checks:
+	@echo "🏥 Running provider health checks..."
+	uv run --group dev python scripts/health_checks.py
+
 test-run-tutorials:
 	@echo "🧪 Running tutorials as e2e tests..."
 	@TUTORIAL_WORKDIR=$$(mktemp -d); \
@@ -582,6 +587,7 @@ clean-test-coverage:
         install install-dev install-dev-notebooks install-dev-recipes \
         lint lint-config lint-engine lint-fix lint-fix-config lint-fix-engine lint-fix-interface lint-interface \
         perf-import publish serve-docs-locally show-versions \
+        health-checks \
         test test-config test-config-isolated test-e2e test-engine test-engine-isolated \
         test-interface test-interface-isolated test-isolated \
         test-run-all-examples test-run-recipes test-run-tutorials \
diff --git a/README.md b/README.md
@@ -2,7 +2,7 @@
 
 [![CI](https://github.com/NVIDIA-NeMo/DataDesigner/actions/workflows/ci.yml/badge.svg)](https://github.com/NVIDIA-NeMo/DataDesigner/actions/workflows/ci.yml)
 [![License](https://img.shields.io/badge/License-Apache_2.0-blue.svg)](https://opensource.org/licenses/Apache-2.0)
-[![Python 3.10 - 3.13](https://img.shields.io/badge/🐍_Python-3.10_|_3.11_|_3.12_|_3.13-blue.svg)](https://www.python.org/downloads/) [![NeMo Microservices](https://img.shields.io/badge/NeMo-Microservices-76b900)](https://docs.nvidia.com/nemo/microservices/latest/index.html) [![Code](https://img.shields.io/badge/Code-Documentation-8A2BE2.svg)](https://nvidia-nemo.github.io/DataDesigner/)
+[![Python 3.10 - 3.13](https://img.shields.io/badge/🐍_Python-3.10_|_3.11_|_3.12_|_3.13-blue.svg)](https://www.python.org/downloads/) [![NeMo Microservices](https://img.shields.io/badge/NeMo-Microservices-76b900)](https://docs.nvidia.com/nemo/microservices/latest/index.html) [![Code](https://img.shields.io/badge/Code-Documentation-8A2BE2.svg)](https://nvidia-nemo.github.io/DataDesigner/) ![Tokens](https://img.shields.io/badge/100+_Billion-Tokens_Generated-76b900.svg?logo=nvidia&logoColor=white)
 
 **Generate high-quality synthetic datasets from scratch or using your own seed data.**
 
@@ -148,11 +148,11 @@ To disable telemetry capture, set `NEMO_TELEMETRY_ENABLED=false`.
 
 ### Top Models
 
-This chart represents the breakdown of models used for Data Designer across all synthetic data generation jobs from 12/18/2025 to 1/14/2026.
+This chart represents the breakdown of models used for Data Designer across all synthetic data generation jobs from 1/5/2026 to 2/5/2026.
 
 ![Top models used for synthetic data generation](docs/images/top-models.png)
 
-_Last updated on 1/14/2026_
+_Last updated on 2/05/2026_
 
 ---
 
diff --git a/docs/images/top-models.png b/docs/images/top-models.png
diff --git a/packages/data-designer-config/src/data_designer/config/__init__.py b/packages/data-designer-config/src/data_designer/config/__init__.py
@@ -99,6 +99,7 @@
         RemoteValidatorParams,
         ValidatorType,
     )
+    from data_designer.config.version import get_library_version  # noqa: F401
 
 # Base module path and submodule paths for lazy imports
 _MOD_BASE = "data_designer.config"
@@ -202,6 +203,8 @@
     "LocalCallableValidatorParams": (_MOD_VALIDATOR_PARAMS, "LocalCallableValidatorParams"),
     "RemoteValidatorParams": (_MOD_VALIDATOR_PARAMS, "RemoteValidatorParams"),
     "ValidatorType": (_MOD_VALIDATOR_PARAMS, "ValidatorType"),
+    # version
+    "get_library_version": (f"{_MOD_BASE}.version", "get_library_version"),
 }
 
 __all__ = list(_LAZY_IMPORTS.keys())
diff --git a/packages/data-designer-config/src/data_designer/config/config_builder.py b/packages/data-designer-config/src/data_designer/config/config_builder.py
@@ -7,6 +7,7 @@
 import logging
 from pathlib import Path
 
+from pydantic import computed_field
 from pygments import highlight
 from pygments.formatters import HtmlFormatter
 from pygments.lexers import PythonLexer
@@ -47,6 +48,7 @@
 from data_designer.config.utils.io_helpers import serialize_data, smart_load_yaml
 from data_designer.config.utils.misc import can_run_data_designer_locally, json_indent_list_of_strings, kebab_to_snake
 from data_designer.config.utils.type_helpers import resolve_string_enum
+from data_designer.config.version import get_library_version
 
 logger = logging.getLogger(__name__)
 
@@ -64,6 +66,10 @@ class BuilderConfig(ExportableConfigBase):
 
     data_designer: DataDesignerConfig
 
+    @computed_field
+    def library_version(self) -> str:
+        return get_library_version()
+
 
 class DataDesignerConfigBuilder:
     """Config builder for Data Designer configurations.
diff --git a/packages/data-designer-config/src/data_designer/config/version.py b/packages/data-designer-config/src/data_designer/config/version.py
@@ -0,0 +1,14 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+from __future__ import annotations
+
+import importlib.metadata
+
+
+def get_library_version() -> str:
+    """Get the installed library version, or 'unknown' if not available."""
+    try:
+        return importlib.metadata.version("data-designer-config")
+    except importlib.metadata.PackageNotFoundError:
+        return "unknown"
diff --git a/packages/data-designer-config/tests/config/test_config_builder.py b/packages/data-designer-config/tests/config/test_config_builder.py
@@ -41,6 +41,7 @@
 from data_designer.config.utils.code_lang import CodeLang
 from data_designer.config.utils.info import ConfigBuilderInfo
 from data_designer.config.validator_params import CodeValidatorParams
+from data_designer.config.version import get_library_version
 from data_designer.lazy_heavy_imports import pd
 
 if TYPE_CHECKING:
@@ -360,6 +361,17 @@ def test_add_profiler(stub_empty_builder):
         stub_empty_builder.add_profiler("invalid")
 
 
+def test_builder_config_library_version(stub_data_designer_builder):
+    builder_config = stub_data_designer_builder.get_builder_config()
+    assert isinstance(builder_config.library_version, str)
+    assert builder_config.library_version == get_library_version()
+
+    # Verify it is included in serialization
+    dumped = builder_config.model_dump()
+    assert "library_version" in dumped
+    assert dumped["library_version"] == builder_config.library_version
+
+
 def test_build(stub_data_designer_builder):
     # verify transformation to config object
     ndd_config = stub_data_designer_builder.build()
diff --git a/packages/data-designer-config/tests/config/test_version.py b/packages/data-designer-config/tests/config/test_version.py
@@ -0,0 +1,24 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+from __future__ import annotations
+
+import importlib.metadata
+from unittest.mock import patch
+
+from data_designer.config.version import get_library_version
+
+
+def test_get_library_version_returns_valid_string() -> None:
+    version = get_library_version()
+    assert isinstance(version, str)
+    assert version != "unknown"
+    assert len(version) > 0
+
+
+def test_get_library_version_returns_unknown_on_missing_package() -> None:
+    with patch(
+        "data_designer.config.version.importlib.metadata.version",
+        side_effect=importlib.metadata.PackageNotFoundError("data-designer-config"),
+    ):
+        assert get_library_version() == "unknown"
diff --git a/packages/data-designer-engine/src/data_designer/engine/dataset_builders/column_wise_builder.py b/packages/data-designer-engine/src/data_designer/engine/dataset_builders/column_wise_builder.py
@@ -4,7 +4,6 @@
 from __future__ import annotations
 
 import functools
-import importlib.metadata
 import logging
 import time
 import uuid
@@ -21,6 +20,7 @@
     ProcessorConfig,
     ProcessorType,
 )
+from data_designer.config.version import get_library_version
 from data_designer.engine.column_generators.generators.base import (
     ColumnGenerator,
     ColumnGeneratorWithModel,
@@ -50,7 +50,7 @@
 
 logger = logging.getLogger(__name__)
 
-_CLIENT_VERSION: str = importlib.metadata.version("data-designer-engine")
+_CLIENT_VERSION: str = get_library_version()
 
 
 class ColumnWiseDatasetBuilder:
diff --git a/packages/data-designer/src/data_designer/integrations/huggingface/dataset_card_template.md b/packages/data-designer/src/data_designer/integrations/huggingface/dataset_card_template.md
@@ -86,6 +86,8 @@ Generated with {{ num_columns_configured }} column configuration(s):
 
 ## 📚 Citation
 
+If you use Data Designer in your work, please cite the project as follows:
+
 ```bibtex
 @misc{nemo-data-designer,
   author = {The NeMo Data Designer Team, NVIDIA},
diff --git a/scripts/health_checks.py b/scripts/health_checks.py
@@ -0,0 +1,116 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+"""Health checks for all predefined model providers.
+
+Verifies that each model in each provider can respond to a basic request.
+Providers without an API key set in the environment are skipped.
+
+Usage:
+    uv run python scripts/health_checks.py
+"""
+
+import os
+import sys
+import traceback
+
+from data_designer.config.models import (
+    ChatCompletionInferenceParams,
+    EmbeddingInferenceParams,
+    ModelConfig,
+    ModelProvider,
+)
+from data_designer.config.utils.constants import (
+    NVIDIA_API_KEY_ENV_VAR_NAME,
+    NVIDIA_PROVIDER_NAME,
+    OPENAI_API_KEY_ENV_VAR_NAME,
+    OPENAI_PROVIDER_NAME,
+    OPENROUTER_API_KEY_ENV_VAR_NAME,
+    OPENROUTER_PROVIDER_NAME,
+    PREDEFINED_PROVIDERS,
+    PREDEFINED_PROVIDERS_MODEL_MAP,
+)
+from data_designer.engine.model_provider import ModelProviderRegistry
+from data_designer.engine.models.facade import ModelFacade
+from data_designer.engine.secret_resolver import EnvironmentResolver
+
+PROVIDER_API_KEY_ENV_VARS = {
+    NVIDIA_PROVIDER_NAME: NVIDIA_API_KEY_ENV_VAR_NAME,
+    OPENAI_PROVIDER_NAME: OPENAI_API_KEY_ENV_VAR_NAME,
+    OPENROUTER_PROVIDER_NAME: OPENROUTER_API_KEY_ENV_VAR_NAME,
+}
+
+
+def _get_provider_registry(provider_name: str) -> ModelProviderRegistry:
+    provider_data = next(p for p in PREDEFINED_PROVIDERS if p["name"] == provider_name)
+    provider = ModelProvider(**provider_data)
+    return ModelProviderRegistry(providers=[provider])
+
+
+def _check_model(provider_name: str, model_type: str) -> None:
+    provider_registry = _get_provider_registry(provider_name)
+    secret_resolver = EnvironmentResolver()
+
+    model_info = PREDEFINED_PROVIDERS_MODEL_MAP[provider_name][model_type]
+    model_name = model_info["model"]
+    inference_params = model_info["inference_parameters"]
+
+    if model_type == "embedding":
+        params = EmbeddingInferenceParams(**inference_params)
+    else:
+        params = ChatCompletionInferenceParams(**inference_params)
+
+    model_config = ModelConfig(
+        alias=f"{provider_name}-{model_type}",
+        model=model_name,
+        inference_parameters=params,
+        provider=provider_name,
+    )
+
+    facade = ModelFacade(model_config, secret_resolver, provider_registry)
+
+    if model_type == "embedding":
+        result = facade.generate_text_embeddings(
+            input_texts=["Hello!"],
+            skip_usage_tracking=True,
+        )
+        assert len(result) == 1 and len(result[0]) > 0
+    else:
+        result, _ = facade.generate(
+            prompt="Say 'OK' and nothing else.",
+            parser=lambda x: x,
+            system_prompt="You are a helpful assistant.",
+            max_correction_steps=0,
+            max_conversation_restarts=0,
+            skip_usage_tracking=True,
+        )
+        assert isinstance(result, str) and len(result) > 0
+
+
+def main() -> int:
+    passed, failed, skipped = 0, 0, 0
+
+    for provider_name, env_var in PROVIDER_API_KEY_ENV_VARS.items():
+        if not os.environ.get(env_var):
+            models = list(PREDEFINED_PROVIDERS_MODEL_MAP[provider_name])
+            skipped += len(models)
+            print(f"SKIP  {provider_name} ({env_var} not set)")
+            continue
+
+        for model_type in PREDEFINED_PROVIDERS_MODEL_MAP[provider_name]:
+            label = f"{provider_name}/{model_type}"
+            try:
+                _check_model(provider_name, model_type)
+                passed += 1
+                print(f"PASS  {label}")
+            except Exception:
+                failed += 1
+                tb = traceback.format_exc()
+                print(f"FAIL  {label}\n{tb}")
+
+    print(f"\n{passed} passed, {failed} failed, {skipped} skipped")
+    return 1 if failed else 0
+
+
+if __name__ == "__main__":
+    sys.exit(main())

Original file line number	Diff line number	Diff line change
`@@ -99,6 +99,7 @@`
`99`	`99`	`RemoteValidatorParams,`
`100`	`100`	`ValidatorType,`
`101`	`101`	`)`
	`102`	`+ from data_designer.config.version import get_library_version # noqa: F401`
`102`	`103`
`103`	`104`	`# Base module path and submodule paths for lazy imports`
`104`	`105`	`_MOD_BASE = "data_designer.config"`
`@@ -202,6 +203,8 @@`
`202`	`203`	`"LocalCallableValidatorParams": (_MOD_VALIDATOR_PARAMS, "LocalCallableValidatorParams"),`
`203`	`204`	`"RemoteValidatorParams": (_MOD_VALIDATOR_PARAMS, "RemoteValidatorParams"),`
`204`	`205`	`"ValidatorType": (_MOD_VALIDATOR_PARAMS, "ValidatorType"),`
	`206`	`+ # version`
	`207`	`+ "get_library_version": (f"{_MOD_BASE}.version", "get_library_version"),`
`205`	`208`	`}`
`206`	`209`
`207`	`210`	`__all__ = list(_LAZY_IMPORTS.keys())`