feat: check_models external readiness check (#712)

Author: mikeknep

packages/data-designer-engine/src/data_designer/engine/dataset_builders/dataset_builder.py

@@ -28,12 +28,12 @@
 from data_designer.config.utils.type_helpers import StrEnum
 from data_designer.config.utils.warning_helpers import warn_at_caller
 from data_designer.config.version import get_library_version
+from data_designer.engine import flags
 from data_designer.engine.column_generators.generators.base import (
     ColumnGenerator,
     ColumnGeneratorWithModel,
     GenerationStrategy,
 )
-from data_designer.engine.column_generators.utils.generator_classification import column_type_is_model_generated
 from data_designer.engine.compiler import compile_data_designer_config
 from data_designer.engine.context import current_row_group, current_row_group_start_offset
 from data_designer.engine.dataset_builders.errors import DatasetGenerationError
@@ -59,9 +59,11 @@
     strip_skip_metadata_from_records,
 )
 from data_designer.engine.dataset_builders.utils.sticky_progress_bar import StickyProgressBar
+from data_designer.engine.models.clients.adapters.http_model_client import ClientConcurrencyMode
 from data_designer.engine.models.telemetry import InferenceEvent, NemoSourceEnum, TaskStatusEnum, TelemetryHandler
 from data_designer.engine.processing.processors.base import Processor
 from data_designer.engine.processing.processors.drop_columns import DropColumnsProcessor
+from data_designer.engine.readiness import run_readiness_check
 from data_designer.engine.registry.data_designer_registry import DataDesignerRegistry
 from data_designer.engine.resources.resource_provider import ResourceProvider
 from data_designer.engine.storage.artifact_storage import (
@@ -82,12 +84,12 @@
 
 logger = logging.getLogger(__name__)
 
-# Async engine is the default execution path. Set ``DATA_DESIGNER_ASYNC_ENGINE=0``
-# to opt back into the legacy sync engine for one transitional release; the sync
-# path is scheduled for removal afterwards.
-DATA_DESIGNER_ASYNC_ENGINE = os.environ.get("DATA_DESIGNER_ASYNC_ENGINE", "1") == "1"
+# The async-engine flag now lives in ``data_designer.engine.flags`` so the
+# engine, the public interface, and the readiness module can share one source
+# of truth. Always read ``flags.DATA_DESIGNER_ASYNC_ENGINE`` rather than caching
+# a local copy so monkeypatches in tests are visible.
 
-if DATA_DESIGNER_ASYNC_ENGINE:
+if flags.DATA_DESIGNER_ASYNC_ENGINE:
     import asyncio
 
     from data_designer.engine.dataset_builders.async_scheduler import (
@@ -193,7 +195,7 @@ def __init__(
         self._task_traces: list[TaskTrace] = []
         self._registry = registry or DataDesignerRegistry()
         self._graph: ExecutionGraph | None = None
-        self._use_async: bool = DATA_DESIGNER_ASYNC_ENGINE
+        self._use_async: bool = flags.DATA_DESIGNER_ASYNC_ENGINE
         # Structured signal: set by _build_async if the scheduler hit early shutdown.
         # Stays at defaults for sync-engine and successful async runs. Reset at
         # the start of each public run path so reused builder instances don't
@@ -275,10 +277,6 @@ def single_column_configs(self) -> list[ColumnConfigT]:
     def single_column_config_by_name(self) -> dict[str, ColumnConfigT]:
         return {config.name: config for config in self.single_column_configs}
 
-    @functools.cached_property
-    def llm_generated_column_configs(self) -> list[ColumnConfigT]:
-        return [config for config in self.single_column_configs if column_type_is_model_generated(config.column_type)]
-
     def build(
         self,
         *,
@@ -314,9 +312,13 @@ def build(
             Path to the generated dataset directory.
         """
         self._reset_run_state()
+        self._use_async = flags.DATA_DESIGNER_ASYNC_ENGINE and self._resolve_async_compatibility()
 
-        self._run_model_health_check_if_needed()
-        self._run_mcp_tool_check_if_needed()
+        run_readiness_check(
+            self.single_column_configs,
+            self._resource_provider,
+            client_concurrency_mode=ClientConcurrencyMode.ASYNC if self._use_async else ClientConcurrencyMode.SYNC,
+        )
 
         # For IF_POSSIBLE and ALWAYS: check config compatibility before touching the artifact
         # directory. _check_resume_config_compatibility() must NOT access base_dataset_path
@@ -386,7 +388,6 @@ def build(
                 "start a new generation run."
             )
 
-        self._use_async = DATA_DESIGNER_ASYNC_ENGINE and self._resolve_async_compatibility()
         if self._use_async:
             self._build_async(generators, num_records, buffer_size, on_batch_complete, resume=resume)
         elif resume == ResumeMode.ALWAYS:
@@ -657,8 +658,12 @@ def _build_with_resume(
 
     def build_preview(self, *, num_records: int) -> pd.DataFrame:
         self._reset_run_state()
-        self._run_model_health_check_if_needed()
-        self._run_mcp_tool_check_if_needed()
+        self._use_async = flags.DATA_DESIGNER_ASYNC_ENGINE and self._resolve_async_compatibility()
+        run_readiness_check(
+            self.single_column_configs,
+            self._resource_provider,
+            client_concurrency_mode=ClientConcurrencyMode.ASYNC if self._use_async else ClientConcurrencyMode.SYNC,
+        )
 
         # Set media storage to DATAFRAME mode for preview - base64 stored directly in DataFrame
         if self._has_image_columns():
@@ -667,7 +672,6 @@ def build_preview(self, *, num_records: int) -> pd.DataFrame:
         generators, self._graph = self._initialize_generators_and_graph()
         start_time = time.perf_counter()
 
-        self._use_async = DATA_DESIGNER_ASYNC_ENGINE and self._resolve_async_compatibility()
         if self._use_async:
             dataset = self._build_async_preview(generators, num_records)
         else:
@@ -1407,38 +1411,6 @@ def _merge_skipped_and_generated(
             batch.append(gen_result)
         return batch
 
-    def _run_model_health_check_if_needed(self) -> None:
-        model_aliases: set[str] = set()
-        for config in self.single_column_configs:
-            model_aliases.update(config.get_model_aliases())
-
-        if not model_aliases:
-            return
-
-        if DATA_DESIGNER_ASYNC_ENGINE:
-            loop = ensure_async_engine_loop()
-            future = asyncio.run_coroutine_threadsafe(
-                self._resource_provider.model_registry.arun_health_check(list(model_aliases)),
-                loop,
-            )
-            try:
-                future.result(timeout=180)
-            except TimeoutError:
-                future.cancel()
-                raise
-        else:
-            self._resource_provider.model_registry.run_health_check(list(model_aliases))
-
-    def _run_mcp_tool_check_if_needed(self) -> None:
-        tool_aliases = sorted(
-            {config.tool_alias for config in self.llm_generated_column_configs if getattr(config, "tool_alias", None)}
-        )
-        if not tool_aliases:
-            return
-        if self._resource_provider.mcp_registry is None:
-            raise DatasetGenerationError(f"Tool alias(es) {tool_aliases!r} specified but no MCPRegistry configured.")
-        self._resource_provider.mcp_registry.run_health_check(tool_aliases)
-
     def _setup_fan_out(
         self,
         generator: ColumnGeneratorWithModelRegistry,

packages/data-designer-engine/src/data_designer/engine/flags.py

@@ -0,0 +1,19 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+"""Engine-wide feature flags read from environment variables.
+
+This module exists so the engine, the public interface, and the readiness
+module can share a single source of truth for runtime mode flags without
+forming an import cycle. Tests patch values here to flip behavior for a
+single test scope.
+"""
+
+from __future__ import annotations
+
+import os
+
+# Async engine is the default execution path. Set ``DATA_DESIGNER_ASYNC_ENGINE=0``
+# to opt back into the legacy sync engine for one transitional release; the sync
+# path is scheduled for removal afterwards.
+DATA_DESIGNER_ASYNC_ENGINE: bool = os.environ.get("DATA_DESIGNER_ASYNC_ENGINE", "1") == "1"

packages/data-designer-engine/src/data_designer/engine/models/registry.py

@@ -8,6 +8,8 @@
 
 from data_designer.config.models import GenerationType, ModelConfig
 from data_designer.engine.model_provider import ModelProvider, ModelProviderRegistry
+from data_designer.engine.models.errors import ModelGenerationValidationFailureError
+from data_designer.engine.models.parsers.errors import ParserException
 from data_designer.engine.models.usage import ModelUsageStats, RequestUsageStats, TokenCountSource, TokenUsageStats
 from data_designer.engine.secret_resolver import SecretResolver
 from data_designer.logging import LOG_INDENT
@@ -27,6 +29,21 @@
 logger = logging.getLogger(__name__)
 
 
+def _parse_health_check_chat_response(response: str) -> str:
+    if not isinstance(response, str) or not response:
+        raise ParserException("Health check response must be non-empty text.")
+    return response
+
+
+def _validate_health_check_embedding_response(vectors: list[list[float]], *, model_alias: str) -> None:
+    if not isinstance(vectors, list) or len(vectors) != 1 or not isinstance(vectors[0], list) or not vectors[0]:
+        raise ModelGenerationValidationFailureError(
+            f"Health check for model alias {model_alias!r} returned an invalid embedding response.",
+            detail="Expected exactly one non-empty embedding vector.",
+            failure_kind="validation_error",
+        )
+
+
 def format_reasoning_token_count(reasoning_token_count: int, source: TokenCountSource | str | None) -> str:
     if source == TokenCountSource.ESTIMATED or source == TokenCountSource.ESTIMATED.value:
         return f"{reasoning_token_count} (estimated)"
@@ -241,15 +258,16 @@ def run_health_check(self, model_aliases: list[str]) -> None:
             )
             try:
                 if model.model_generation_type == GenerationType.EMBEDDING:
-                    model.generate_text_embeddings(
+                    vectors = model.generate_text_embeddings(
                         input_texts=["Hello!"],
                         skip_usage_tracking=True,
                         purpose="running health checks",
                     )
+                    _validate_health_check_embedding_response(vectors, model_alias=model_alias)
                 elif model.model_generation_type == GenerationType.CHAT_COMPLETION:
                     model.generate(
                         prompt="Hello!",
-                        parser=lambda x: x,
+                        parser=_parse_health_check_chat_response,
                         system_prompt="You are a helpful assistant.",
                         max_correction_steps=0,
                         max_conversation_restarts=0,
@@ -286,15 +304,16 @@ async def arun_health_check(self, model_aliases: list[str]) -> None:
             )
             try:
                 if model.model_generation_type == GenerationType.EMBEDDING:
-                    await model.agenerate_text_embeddings(
+                    vectors = await model.agenerate_text_embeddings(
                         input_texts=["Hello!"],
                         skip_usage_tracking=True,
                         purpose="running health checks",
                     )
+                    _validate_health_check_embedding_response(vectors, model_alias=model_alias)
                 elif model.model_generation_type == GenerationType.CHAT_COMPLETION:
                     await model.agenerate(
                         prompt="Hello!",
-                        parser=lambda x: x,
+                        parser=_parse_health_check_chat_response,
                         system_prompt="You are a helpful assistant.",
                         max_correction_steps=0,
                         max_conversation_restarts=0,

packages/data-designer-engine/src/data_designer/engine/readiness.py

@@ -0,0 +1,127 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+"""External-readiness checks for a DataDesigner workload.
+
+A "readiness" check is a pre-flight probe of every external resource a
+configuration depends on: each referenced model alias is sent a tiny
+generation request, and every referenced MCP tool alias is contacted to
+confirm its server is reachable.
+
+This module hosts the shared logic invoked from two places:
+
+- ``DatasetBuilder.build`` / ``DatasetBuilder.build_preview`` — at the start
+  of a workload, to fail fast before any expensive work begins.
+- ``DataDesigner.check_models`` — exposed publicly so users can verify
+  external dependencies are responsive without triggering a workload.
+
+The two callers must use the same code path here so the standalone method
+cannot drift from the workload-startup gate.
+"""
+
+from __future__ import annotations
+
+import logging
+from collections.abc import Sequence
+from typing import TYPE_CHECKING
+
+from data_designer.engine.column_generators.utils.generator_classification import column_type_is_model_generated
+from data_designer.engine.dataset_builders.errors import DatasetGenerationError
+from data_designer.engine.models.clients.adapters.http_model_client import ClientConcurrencyMode
+
+if TYPE_CHECKING:
+    from data_designer.config.column_types import ColumnConfigT
+    from data_designer.engine.resources.resource_provider import ResourceProvider
+
+logger = logging.getLogger(__name__)
+
+# Match the timeout the dataset builder's startup gate has always used.
+_MODEL_HEALTH_CHECK_TIMEOUT_SECONDS = 180
+
+
+def run_readiness_check(
+    column_configs: Sequence[ColumnConfigT],
+    resource_provider: ResourceProvider,
+    *,
+    client_concurrency_mode: ClientConcurrencyMode,
+) -> None:
+    """Probe every model and MCP tool referenced by ``column_configs``.
+
+    For each unique model alias collected from the column configs,
+    ``ModelRegistry.run_health_check`` (or ``arun_health_check`` when async
+    mode is selected) sends a tiny ``"Hello!"`` generation. Models whose ``ModelConfig``
+    has ``skip_health_check=True`` are skipped by the registry. After the
+    model pass, every unique MCP tool alias is probed via
+    ``MCPRegistry.run_health_check``.
+
+    Args:
+        column_configs: The column configs whose ``get_model_aliases()`` and
+            ``tool_alias`` fields determine which aliases are probed.
+        resource_provider: Provides access to the model registry and MCP
+            registry. ``mcp_registry`` may be ``None`` only if no tool
+            aliases are referenced.
+        client_concurrency_mode: Resolved client mode for this run.
+
+    Raises:
+        Typed model errors from ``data_designer.engine.models.errors`` for
+            any failing model probe.
+        DatasetGenerationError: If a tool alias is referenced but no MCP
+            registry is configured on the resource provider.
+        TimeoutError: If async health-check execution exceeds
+            ``_MODEL_HEALTH_CHECK_TIMEOUT_SECONDS``.
+    """
+    _run_model_health_check(column_configs, resource_provider, client_concurrency_mode=client_concurrency_mode)
+    _run_mcp_tool_health_check(column_configs, resource_provider)
+
+
+def _run_model_health_check(
+    column_configs: Sequence[ColumnConfigT],
+    resource_provider: ResourceProvider,
+    *,
+    client_concurrency_mode: ClientConcurrencyMode,
+) -> None:
+    model_aliases: set[str] = set()
+    for config in column_configs:
+        model_aliases.update(config.get_model_aliases())
+
+    if not model_aliases:
+        return
+
+    if client_concurrency_mode == ClientConcurrencyMode.ASYNC:
+        # Defer the async-engine imports to here so users on the legacy sync
+        # engine never pay the import cost.
+        import asyncio
+
+        from data_designer.engine.dataset_builders.utils.async_concurrency import ensure_async_engine_loop
+
+        loop = ensure_async_engine_loop()
+        future = asyncio.run_coroutine_threadsafe(
+            resource_provider.model_registry.arun_health_check(list(model_aliases)),
+            loop,
+        )
+        try:
+            future.result(timeout=_MODEL_HEALTH_CHECK_TIMEOUT_SECONDS)
+        except TimeoutError:
+            future.cancel()
+            raise
+    else:
+        resource_provider.model_registry.run_health_check(list(model_aliases))
+
+
+def _run_mcp_tool_health_check(
+    column_configs: Sequence[ColumnConfigT],
+    resource_provider: ResourceProvider,
+) -> None:
+    # Tool aliases are only meaningful on model-generated column configs.
+    tool_aliases = sorted(
+        {
+            config.tool_alias
+            for config in column_configs
+            if column_type_is_model_generated(config.column_type) and getattr(config, "tool_alias", None)
+        }
+    )
+    if not tool_aliases:
+        return
+    if resource_provider.mcp_registry is None:
+        raise DatasetGenerationError(f"Tool alias(es) {tool_aliases!r} specified but no MCPRegistry configured.")
+    resource_provider.mcp_registry.run_health_check(tool_aliases)

packages/data-designer-engine/src/data_designer/engine/resources/resource_provider.py

@@ -3,7 +3,6 @@
 
 from __future__ import annotations
 
-import os
 from typing import TYPE_CHECKING
 
 from data_designer.config.base import ConfigBase
@@ -13,6 +12,7 @@
 from data_designer.config.run_config import RunConfig
 from data_designer.config.seed_source import SeedSource
 from data_designer.config.utils.type_helpers import StrEnum
+from data_designer.engine import flags
 from data_designer.engine.mcp.factory import create_mcp_registry
 from data_designer.engine.mcp.registry import MCPRegistry
 from data_designer.engine.model_provider import (
@@ -148,9 +148,7 @@ def create_resource_provider(
     # default for backward compatibility.
     if client_concurrency_mode is None:
         client_concurrency_mode = (
-            ClientConcurrencyMode.ASYNC
-            if os.environ.get("DATA_DESIGNER_ASYNC_ENGINE", "1") == "1"
-            else ClientConcurrencyMode.SYNC
+            ClientConcurrencyMode.ASYNC if flags.DATA_DESIGNER_ASYNC_ENGINE else ClientConcurrencyMode.SYNC
         )
 
     effective_run_config = run_config or RunConfig()