align readiness with client mode

Author: mikeknep

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

@@ -59,6 +59,7 @@
     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
@@ -311,8 +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()
 
-        run_readiness_check(self.single_column_configs, self._resource_provider)
+        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
@@ -382,7 +388,6 @@ def build(
                 "start a new generation run."
             )
 
-        self._use_async = flags.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:
@@ -653,7 +658,12 @@ def _build_with_resume(
 
     def build_preview(self, *, num_records: int) -> pd.DataFrame:
         self._reset_run_state()
-        run_readiness_check(self.single_column_configs, self._resource_provider)
+        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():
@@ -662,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 = flags.DATA_DESIGNER_ASYNC_ENGINE and self._resolve_async_compatibility()
         if self._use_async:
             dataset = self._build_async_preview(generators, num_records)
         else:

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

@@ -25,9 +25,9 @@
 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
 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
@@ -42,12 +42,14 @@
 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`` on the async
-    engine) sends a tiny ``"Hello!"`` generation. Models whose ``ModelConfig``
+    ``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``.
@@ -58,6 +60,7 @@ def run_readiness_check(
         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
@@ -67,13 +70,15 @@ def run_readiness_check(
         TimeoutError: If async health-check execution exceeds
             ``_MODEL_HEALTH_CHECK_TIMEOUT_SECONDS``.
     """
-    _run_model_health_check(column_configs, resource_provider)
+    _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:
@@ -82,10 +87,9 @@ def _run_model_health_check(
     if not model_aliases:
         return
 
-    if flags.DATA_DESIGNER_ASYNC_ENGINE:
+    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. Mirrors the gating in
-        # ``dataset_builders.dataset_builder``.
+        # engine never pay the import cost.
         import asyncio
 
         from data_designer.engine.dataset_builders.utils.async_concurrency import ensure_async_engine_loop

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

@@ -17,7 +17,9 @@
 from data_designer.engine import flags
 from data_designer.engine.dataset_builders.errors import DatasetGenerationError
 from data_designer.engine.mcp.registry import MCPRegistry
+from data_designer.engine.models.clients.adapters.http_model_client import ClientConcurrencyMode
 from data_designer.engine.readiness import run_readiness_check
+from data_designer.engine.resources.resource_provider import ResourceProvider
 
 
 @pytest.fixture(autouse=True)
@@ -52,6 +54,14 @@ def _build_columns(
     return builder.build().columns
 
 
+def _run_sync_readiness(column_configs: Sequence[ColumnConfigT], resource_provider: ResourceProvider) -> None:
+    run_readiness_check(
+        column_configs,
+        resource_provider,
+        client_concurrency_mode=ClientConcurrencyMode.SYNC,
+    )
+
+
 # ---------------------------------------------------------------------------
 # Model health check
 # ---------------------------------------------------------------------------
@@ -82,7 +92,7 @@ def _gen_with_two_models(row, generator_params, models):
 
     builder.add_column(CustomColumnConfig(name="custom_col", generator_function=_gen_with_two_models))
 
-    run_readiness_check(builder.build().columns, stub_resource_provider)
+    _run_sync_readiness(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
@@ -99,7 +109,7 @@ def test_run_readiness_check_skips_model_probe_when_no_aliases(
 
     columns = _build_columns(model_configs=stub_model_configs, llm_columns=[])
 
-    run_readiness_check(columns, stub_resource_provider)
+    _run_sync_readiness(columns, stub_resource_provider)
 
     stub_resource_provider.model_registry.run_health_check.assert_not_called()
 
@@ -117,7 +127,7 @@ def test_run_readiness_check_propagates_model_probe_error(
     columns = _build_columns(model_configs=stub_model_configs, llm_columns=[("col", "stub-text")])
 
     with pytest.raises(ModelAuthenticationError, match="bad creds"):
-        run_readiness_check(columns, stub_resource_provider)
+        _run_sync_readiness(columns, stub_resource_provider)
 
 
 # ---------------------------------------------------------------------------
@@ -142,7 +152,7 @@ def test_run_readiness_check_collects_unique_sorted_tool_aliases(
         LLMTextColumnConfig(name="c", prompt="x", model_alias="stub-text", tool_alias="alpha")  # duplicate
     )
 
-    run_readiness_check(builder.build().columns, stub_resource_provider)
+    _run_sync_readiness(builder.build().columns, stub_resource_provider)
 
     mock_mcp_registry.run_health_check.assert_called_once_with(["alpha", "zebra"])
 
@@ -158,7 +168,7 @@ def test_run_readiness_check_skips_tool_probe_when_no_tool_aliases(
 
     columns = _build_columns(model_configs=stub_model_configs, llm_columns=[("col", "stub-text")])
 
-    run_readiness_check(columns, stub_resource_provider)
+    _run_sync_readiness(columns, stub_resource_provider)
 
     mock_mcp_registry.run_health_check.assert_not_called()
 
@@ -176,7 +186,7 @@ def test_run_readiness_check_raises_when_tools_referenced_but_no_mcp_registry(
     builder.add_column(LLMTextColumnConfig(name="col", prompt="x", model_alias="stub-text", tool_alias="missing-tools"))
 
     with pytest.raises(DatasetGenerationError, match="missing-tools"):
-        run_readiness_check(builder.build().columns, stub_resource_provider)
+        _run_sync_readiness(builder.build().columns, stub_resource_provider)
 
 
 def test_run_readiness_check_propagates_tool_probe_error(
@@ -194,7 +204,7 @@ def test_run_readiness_check_propagates_tool_probe_error(
     builder.add_column(LLMTextColumnConfig(name="col", prompt="x", model_alias="stub-text", tool_alias="tools"))
 
     with pytest.raises(RuntimeError, match="mcp down"):
-        run_readiness_check(builder.build().columns, stub_resource_provider)
+        _run_sync_readiness(builder.build().columns, stub_resource_provider)
 
 
 # ---------------------------------------------------------------------------
@@ -218,7 +228,7 @@ def test_run_readiness_check_runs_models_before_tools(
     builder.add_column(LLMTextColumnConfig(name="col", prompt="x", model_alias="stub-text", tool_alias="tools"))
 
     with pytest.raises(ModelAuthenticationError):
-        run_readiness_check(builder.build().columns, stub_resource_provider)
+        _run_sync_readiness(builder.build().columns, stub_resource_provider)
 
     # The MCP probe must not have been reached.
     mock_mcp_registry.run_health_check.assert_not_called()
@@ -235,7 +245,7 @@ def test_run_readiness_check_no_models_no_tools_is_noop(
 
     columns = _build_columns(model_configs=stub_model_configs, llm_columns=[])
 
-    run_readiness_check(columns, stub_resource_provider)
+    _run_sync_readiness(columns, stub_resource_provider)
 
     stub_resource_provider.model_registry.run_health_check.assert_not_called()
     mock_mcp_registry.run_health_check.assert_not_called()
@@ -266,7 +276,7 @@ def test_run_readiness_check_collects_image_model_aliases(
     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)
+    _run_sync_readiness(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
@@ -292,7 +302,7 @@ def test_run_readiness_check_passes_skip_flagged_aliases_to_registry(
         llm_columns=[("col", "stub-text")],
     )
 
-    run_readiness_check(columns, stub_resource_provider)
+    _run_sync_readiness(columns, stub_resource_provider)
 
     stub_resource_provider.model_registry.run_health_check.assert_called_once_with(["stub-text"])
 
@@ -331,7 +341,11 @@ def test_run_readiness_check_dispatches_to_async_registry_under_async_engine(
         patch("data_designer.engine.dataset_builders.utils.async_concurrency.ensure_async_engine_loop"),
         patch("asyncio.run_coroutine_threadsafe", return_value=sentinel_future) as mock_submit,
     ):
-        run_readiness_check(columns, stub_resource_provider)
+        run_readiness_check(
+            columns,
+            stub_resource_provider,
+            client_concurrency_mode=ClientConcurrencyMode.ASYNC,
+        )
 
     # 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"])
@@ -362,6 +376,32 @@ def test_run_readiness_check_cancels_future_and_reraises_on_timeout(
         patch("asyncio.run_coroutine_threadsafe", return_value=sentinel_future),
         pytest.raises(TimeoutError),
     ):
-        run_readiness_check(columns, stub_resource_provider)
+        run_readiness_check(
+            columns,
+            stub_resource_provider,
+            client_concurrency_mode=ClientConcurrencyMode.ASYNC,
+        )
 
     sentinel_future.cancel.assert_called_once()
+
+
+def test_run_readiness_check_uses_sync_registry_for_sync_mode_clients(
+    stub_resource_provider,
+    stub_model_configs,
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    """Readiness follows the explicit client mode, not only the raw async env flag."""
+    monkeypatch.setattr(flags, "DATA_DESIGNER_ASYNC_ENGINE", True)
+    stub_resource_provider.model_registry.run_health_check = Mock()
+    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_readiness_check(columns, stub_resource_provider, client_concurrency_mode=ClientConcurrencyMode.SYNC)
+
+    stub_resource_provider.model_registry.run_health_check.assert_called_once_with(["stub-text"])
+    stub_resource_provider.model_registry.arun_health_check.assert_not_called()

packages/data-designer/src/data_designer/interface/data_designer.py

@@ -578,7 +578,11 @@ def check_models(self, config_builder: DataDesignerConfigBuilder) -> None:
             TimeoutError: If async health-check execution exceeds 180 seconds.
         """
         resource_provider = self._create_resource_provider("check-models", config_builder)
-        run_readiness_check(config_builder.build().columns, resource_provider)
+        run_readiness_check(
+            config_builder.build().columns,
+            resource_provider,
+            client_concurrency_mode=self._resolve_client_concurrency_mode(config_builder),
+        )
 
     def get_default_model_configs(self) -> list[ModelConfig]:
         """Get the default model configurations.

packages/data-designer/tests/interface/test_data_designer.py

@@ -39,6 +39,7 @@
     HuggingFaceSeedSource,
 )
 from data_designer.engine import flags
+from data_designer.engine.models.clients.adapters.http_model_client import ClientConcurrencyMode
 from data_designer.engine.resources.seed_reader import (
     FileSystemSeedReader,
     SeedReaderError,
@@ -1544,9 +1545,33 @@ def test_check_models_invokes_readiness_check(
         data_designer.check_models(config_builder)
 
     assert mock_check.call_count == 1
-    (called_columns, called_resource_provider), _ = mock_check.call_args
+    (called_columns, called_resource_provider), kwargs = mock_check.call_args
     assert [c.name for c in called_columns] == ["text"]
     assert called_resource_provider is not None
+    assert kwargs["client_concurrency_mode"] == ClientConcurrencyMode.ASYNC
+
+
+def test_check_models_passes_sync_mode_for_sync_fallback(
+    stub_artifact_path,
+    stub_model_providers,
+    stub_managed_assets_path,
+    monkeypatch: pytest.MonkeyPatch,
+):
+    """check_models readiness uses the resolved client mode, including allow_resize fallback."""
+    monkeypatch.setattr(flags, "DATA_DESIGNER_ASYNC_ENGINE", True)
+    config_builder = _builder_with_allow_resize()
+    data_designer = DataDesigner(
+        artifact_path=stub_artifact_path,
+        model_providers=stub_model_providers,
+        secret_resolver=PlaintextResolver(),
+        managed_assets_path=stub_managed_assets_path,
+    )
+
+    with patch("data_designer.interface.data_designer.run_readiness_check") as mock_check:
+        data_designer.check_models(config_builder)
+
+    _, kwargs = mock_check.call_args
+    assert kwargs["client_concurrency_mode"] == ClientConcurrencyMode.SYNC
 
 
 def test_check_models_propagates_typed_model_error(