Cleanup

Signed-off-by: Mike Knepper <mknepper@nvidia.com>

Author: mikeknep

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

@@ -34,7 +34,6 @@
     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
@@ -277,10 +276,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,
         *,

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

@@ -22,7 +22,8 @@
 from __future__ import annotations
 
 import logging
-from typing import TYPE_CHECKING, Sequence
+from collections.abc import Sequence
+from typing import TYPE_CHECKING
 
 from data_designer.engine import flags
 from data_designer.engine.column_generators.utils.generator_classification import column_type_is_model_generated
@@ -107,8 +108,7 @@ def _run_mcp_tool_health_check(
     column_configs: Sequence[ColumnConfigT],
     resource_provider: ResourceProvider,
 ) -> None:
-    # Tool aliases are only meaningful on LLM-generated column configs; the
-    # filter here matches the (now-removed) ``DatasetBuilder.llm_generated_column_configs``.
+    # Tool aliases are only meaningful on model-generated column configs.
     tool_aliases = sorted(
         {
             config.tool_alias

packages/data-designer-engine/tests/engine/test_readiness.py

@@ -3,13 +3,16 @@
 
 from __future__ import annotations
 
-from unittest.mock import Mock
+from collections.abc import Sequence
+from unittest.mock import Mock, patch
 
 import pytest
 
 from data_designer.config.column_configs import LLMTextColumnConfig, SamplerColumnConfig
+from data_designer.config.column_types import ColumnConfigT
 from data_designer.config.config_builder import DataDesignerConfigBuilder
 from data_designer.config.custom_column import custom_column_generator
+from data_designer.config.models import ModelConfig
 from data_designer.config.sampler_params import SamplerType, UUIDSamplerParams
 from data_designer.engine import flags
 from data_designer.engine.dataset_builders.errors import DatasetGenerationError
@@ -27,7 +30,12 @@ def _force_sync_engine(monkeypatch: pytest.MonkeyPatch) -> None:
     monkeypatch.setattr(flags, "DATA_DESIGNER_ASYNC_ENGINE", False)
 
 
-def _build_columns(*, model_configs, llm_columns: list[tuple[str, str]] = (), include_sampler: bool = True):
+def _build_columns(
+    *,
+    model_configs: list[ModelConfig],
+    llm_columns: Sequence[tuple[str, str]] = (),
+    include_sampler: bool = True,
+) -> list[ColumnConfigT]:
     """Build a ``DataDesignerConfig`` and return its (already-flat) column configs.
 
     ``llm_columns`` is a list of ``(name, model_alias)`` pairs. ``include_sampler``
@@ -231,3 +239,131 @@ def test_run_readiness_check_no_models_no_tools_is_noop(
 
     stub_resource_provider.model_registry.run_health_check.assert_not_called()
     mock_mcp_registry.run_health_check.assert_not_called()
+
+
+# ---------------------------------------------------------------------------
+# Column-type coverage
+# ---------------------------------------------------------------------------
+
+
+def test_run_readiness_check_collects_image_model_aliases(
+    stub_resource_provider,
+    stub_model_configs,
+) -> None:
+    """Image-generation columns contribute their model aliases like LLM columns do.
+
+    The dataset builder dispatches probes by ``model_generation_type`` inside
+    ``ModelRegistry.run_health_check``; readiness is generation-type-agnostic
+    and must surface every alias regardless of column kind.
+    """
+    from data_designer.config.column_configs import ImageColumnConfig
+
+    stub_resource_provider.model_registry.run_health_check = Mock()
+    stub_resource_provider.mcp_registry = None
+
+    builder = DataDesignerConfigBuilder(model_configs=stub_model_configs)
+    builder.add_column(SamplerColumnConfig(name="seed_id", sampler_type=SamplerType.UUID, params=UUIDSamplerParams()))
+    builder.add_column(LLMTextColumnConfig(name="caption", prompt="x", model_alias="stub-text"))
+    builder.add_column(ImageColumnConfig(name="picture", prompt="y", model_alias="stub-image"))
+
+    run_readiness_check(builder.build().columns, stub_resource_provider)
+
+    stub_resource_provider.model_registry.run_health_check.assert_called_once()
+    (called_aliases,), _ = stub_resource_provider.model_registry.run_health_check.call_args
+    assert set(called_aliases) == {"stub-text", "stub-image"}
+
+
+def test_run_readiness_check_passes_skip_flagged_aliases_to_registry(
+    stub_resource_provider,
+    stub_model_configs,
+) -> None:
+    """Readiness does not pre-filter ``skip_health_check=True`` aliases.
+
+    The skip decision lives in ``ModelRegistry.run_health_check`` (covered by
+    ``test_model_registry``). Readiness's contract is "pass every referenced
+    alias through and let the registry decide" — verified here so future edits
+    don't accidentally start filtering at this layer.
+    """
+    stub_resource_provider.model_registry.run_health_check = Mock()
+    stub_resource_provider.mcp_registry = None
+
+    columns = _build_columns(
+        model_configs=stub_model_configs,
+        llm_columns=[("col", "stub-text")],
+    )
+
+    run_readiness_check(columns, stub_resource_provider)
+
+    stub_resource_provider.model_registry.run_health_check.assert_called_once_with(["stub-text"])
+
+
+# ---------------------------------------------------------------------------
+# Async dispatch
+# ---------------------------------------------------------------------------
+
+
+def test_run_readiness_check_dispatches_to_async_registry_under_async_engine(
+    stub_resource_provider,
+    stub_model_configs,
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    """When the async engine is selected, model probes route through ``arun_health_check``.
+
+    The autouse fixture pins sync; this test overrides for the async path so the
+    branch in ``readiness._run_model_health_check`` gets coverage.
+    """
+    monkeypatch.setattr(flags, "DATA_DESIGNER_ASYNC_ENGINE", True)
+    stub_resource_provider.model_registry.arun_health_check = Mock()
+    stub_resource_provider.mcp_registry = None
+
+    columns = _build_columns(
+        model_configs=stub_model_configs,
+        llm_columns=[("col", "stub-text")],
+    )
+
+    # ``run_coroutine_threadsafe`` returns a Future; we want the readiness wrapper
+    # to call ``.result(timeout=...)`` on it, so install a Mock future whose
+    # ``.result`` returns ``None`` (success).
+    sentinel_future = Mock()
+    sentinel_future.result.return_value = None
+
+    fake_loop = Mock()
+
+    with (
+        patch("data_designer.engine.readiness.ensure_async_engine_loop", return_value=fake_loop, create=True),
+        patch("asyncio.run_coroutine_threadsafe", return_value=sentinel_future) as mock_submit,
+    ):
+        run_readiness_check(columns, stub_resource_provider)
+
+    # The async coroutine was created from arun_health_check and submitted to the loop.
+    stub_resource_provider.model_registry.arun_health_check.assert_called_once_with(["stub-text"])
+    mock_submit.assert_called_once()
+    sentinel_future.result.assert_called_once_with(timeout=180)
+
+
+def test_run_readiness_check_cancels_future_and_reraises_on_timeout(
+    stub_resource_provider,
+    stub_model_configs,
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    """A 180-second timeout cancels the future and re-raises ``TimeoutError``."""
+    monkeypatch.setattr(flags, "DATA_DESIGNER_ASYNC_ENGINE", True)
+    stub_resource_provider.model_registry.arun_health_check = Mock()
+    stub_resource_provider.mcp_registry = None
+
+    columns = _build_columns(
+        model_configs=stub_model_configs,
+        llm_columns=[("col", "stub-text")],
+    )
+
+    sentinel_future = Mock()
+    sentinel_future.result.side_effect = TimeoutError()
+
+    with (
+        patch("data_designer.engine.readiness.ensure_async_engine_loop", return_value=Mock(), create=True),
+        patch("asyncio.run_coroutine_threadsafe", return_value=sentinel_future),
+        pytest.raises(TimeoutError),
+    ):
+        run_readiness_check(columns, stub_resource_provider)
+
+    sentinel_future.cancel.assert_called_once()

packages/data-designer/src/data_designer/cli/commands/check_models.py

@@ -31,6 +31,9 @@ def check_models_command(
         # Check models referenced by a YAML config
         data-designer check-models my_config.yaml
 
+        # Check models referenced by a remote config URL
+        data-designer check-models https://example.com/my_config.yaml
+
         # Check models referenced by a Python module
         data-designer check-models my_config.py
     """

packages/data-designer/src/data_designer/cli/controllers/generation_controller.py

@@ -17,6 +17,7 @@
 from data_designer.config.errors import InvalidConfigError
 from data_designer.config.utils.constants import DEFAULT_DISPLAY_WIDTH
 from data_designer.engine.storage.artifact_storage import ResumeMode
+from data_designer.errors import DataDesignerError
 from data_designer.interface import DataDesigner
 from data_designer.logging import LOG_INDENT
 
@@ -130,6 +131,9 @@ def run_check_models(self, config_source: str) -> None:
         try:
             data_designer = DataDesigner()
             data_designer.check_models(config_builder)
+        except DataDesignerError as e:
+            print_error(f"Model health check failed ({type(e).__name__}): {e}")
+            raise typer.Exit(code=1)
         except Exception as e:
             print_error(f"Model health check failed: {e}")
             raise typer.Exit(code=1)

packages/data-designer/tests/cli/controllers/test_generation_controller.py

@@ -701,7 +701,7 @@ def test_run_check_models_config_load_error(mock_load_config: MagicMock) -> None
 @patch(f"{_CTRL}.DataDesigner")
 @patch(f"{_CTRL}.load_config_builder")
 def test_run_check_models_health_check_failure(mock_load_config: MagicMock, mock_dd_cls: MagicMock) -> None:
-    """check_models exits with code 1 when a probe fails."""
+    """check_models exits with code 1 when a probe fails with a generic exception."""
     mock_load_config.return_value = MagicMock(spec=DataDesignerConfigBuilder)
     mock_dd = MagicMock()
     mock_dd_cls.return_value = mock_dd
@@ -714,6 +714,33 @@ def test_run_check_models_health_check_failure(mock_load_config: MagicMock, mock
     assert exc_info.value.exit_code == 1
 
 
+@patch(f"{_CTRL}.DataDesigner")
+@patch(f"{_CTRL}.load_config_builder")
+def test_run_check_models_typed_error_includes_class_name(
+    mock_load_config: MagicMock, mock_dd_cls: MagicMock, capsys: pytest.CaptureFixture[str]
+) -> None:
+    """Typed engine errors exit 1 and surface the error class name to the user.
+
+    Without this, an authentication failure and a connection failure look identical
+    on the terminal, defeating the purpose of typed engine errors.
+    """
+    from data_designer.engine.models.errors import ModelAuthenticationError
+
+    mock_load_config.return_value = MagicMock(spec=DataDesignerConfigBuilder)
+    mock_dd = MagicMock()
+    mock_dd_cls.return_value = mock_dd
+    mock_dd.check_models.side_effect = ModelAuthenticationError("bad creds")
+
+    controller = GenerationController()
+    with pytest.raises(typer.Exit) as exc_info:
+        controller.run_check_models(config_source="config.yaml")
+
+    assert exc_info.value.exit_code == 1
+    captured = capsys.readouterr()
+    assert "ModelAuthenticationError" in captured.out
+    assert "bad creds" in captured.out
+
+
 # ---------------------------------------------------------------------------
 # run_create tests
 # ---------------------------------------------------------------------------