DataDesigner/packages/data-designer-engine/tests/engine/test_configurable_task.py
Johnny Greco 1439bbea7e
chore: Improve CLI startup with lazy heavy import cleanup (#330)
* perf: defer heavy imports to improve CLI startup time

Move expensive imports (engine, models, controllers) out of the module-level import path so that data-designer --help and other non-generation commands no longer pay the full startup cost.

Key changes:
- Defer controller imports to inside command functions
- Remove eager re-export chains from CLI package __init__ files
- Move default-settings bootstrap into load_config_builder() and DataDesigner.__init__() instead of running at import time
- Add lazy __getattr__ exports in interface/__init__.py
- Replace module-level tokenizer init with cached lazy getter
- Fix ModelProvider import to use config layer instead of engine
- Update test mock paths to match new import locations

Reduces CLI import-time from ~1.67s to ~0.46s.

* perf: defer pandas/numpy in io_helpers and add config_list benchmark

- Replace eager `from lazy_heavy_imports import pd, np` in io_helpers
  with module-level __getattr__ (for backwards-compatible external
  access / test mocks) and function-level imports in the 3 functions
  that actually use them (read_parquet_dataset, smart_load_dataframe,
  _convert_to_serializable). Importing io_helpers no longer triggers
  pandas/numpy loading.
- Defer heavy imports in list and reset CLI commands into function
  bodies to avoid loading repositories, Rich, and prompt_toolkit at
  module import time.
- Add `config_list` (data-designer config list) measurement to the
  CLI startup benchmark with isolated cold measurement in a separate
  venv and a --skip-config-list-check flag.
- Update test mock paths to match new import locations.

* Refine lazy import usage and TYPE_CHECKING cleanup

* Run license header updater on PR-touched files

* fix: update sqlfluff mock target for lazy imports in test_sql

* perf: cache globals() in lazy __getattr__ to avoid repeated lookups

Add globals() caching and explanatory comment to all three lazy
__getattr__ implementations (lazy_heavy_imports, config/__init__,
interface/__init__) so subsequent attribute accesses bypass __getattr__.

* perf: lazy CLI command loading and deferred heavy import evaluations

- Add LazyTyperGroup to defer command module loading until invocation, allowing module-level imports in all CLI command files

- Split DataFrameSeedSource into seed_source_dataframe.py to isolate pandas dependency from other seed source classes

- Move TypeVar/TypeAlias definitions (DataT, NumpyArray1dT, RadomStateT, EngineT) to TYPE_CHECKING blocks with runtime fallbacks

- Wrap module-level constants in lru_cache (phone_number parquet data, jsonschema validator) to defer I/O and heavy imports to first use

- Update test mock targets to patch at usage-site for module-level imports

* refactor: use direct pandas import in seed_source_dataframe

Drop lazy-loading for pandas in DataFrameSeedSource; use direct import
for simplicity.

* update lazy import pattern

* update tests to use lazy import namespace

Switch test modules to import data_designer.lazy_heavy_imports as lazy
and reference heavy libraries through that namespace. This keeps heavy
imports deferred during module import and aligns tests with the new
lazy-import usage pattern.

* tighten import perf test thresholds

Document recent baseline timings and lower the allowed average
import time and timeout so regressions are detected sooner.

* document pandas import requirement

Clarify that Pydantic needs DataFrame resolved at module load and
that keeping the direct import preserves IDE typing support.

* increase timeout time

* use lazy pandas imports in visualization tests

- replace direct pandas usage with lazy.pd in visualization tests to avoid eager imports
- add TYPE_CHECKING pandas import and keep CLI controller imports sorted

* fix lazy pandas runtime usage and preview mocks

Switch sample-record handling to lazy pandas types so runtime paths no longer
depend on TYPE_CHECKING imports. Align preview controller tests to patch the
module-local DataDesigner symbol, preventing real engine invocation in save
results scenarios.
2026-02-18 16:24:15 -05:00

112 lines
3.6 KiB
Python

# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: Apache-2.0
from __future__ import annotations
from unittest.mock import Mock
import pytest
from data_designer.config.base import ConfigBase
from data_designer.engine.configurable_task import ConfigurableTask, DataT, TaskConfigT
from data_designer.engine.models.registry import ModelRegistry
from data_designer.engine.resources.resource_provider import ResourceProvider
from data_designer.engine.storage.artifact_storage import ArtifactStorage
def test_configurable_task_generic_type_variables() -> None:
# DataT constraints are deferred to TYPE_CHECKING to avoid eagerly importing pandas.
# At runtime it is an unconstrained TypeVar; type checkers still see the constraints.
assert DataT.__constraints__ == ()
assert TaskConfigT.__bound__ == ConfigBase
def test_configurable_task_concrete_implementation(tmp_path) -> None:
class TestConfig(ConfigBase):
value: str
class TestTask(ConfigurableTask[TestConfig]):
@classmethod
def get_config_type(cls) -> type[TestConfig]:
return TestConfig
def _validate(self) -> None:
pass
def _initialize(self) -> None:
pass
config = TestConfig(value="test")
artifact_storage = ArtifactStorage(artifact_path=tmp_path)
resource_provider = ResourceProvider(artifact_storage=artifact_storage)
task = TestTask(config=config, resource_provider=resource_provider)
assert task._config == config
assert task._resource_provider == resource_provider
def test_configurable_task_config_validation(tmp_path) -> None:
class TestConfig(ConfigBase):
value: str
class TestTask(ConfigurableTask[TestConfig]):
@classmethod
def get_config_type(cls) -> type[TestConfig]:
return TestConfig
def _validate(self) -> None:
if self._config.value == "invalid":
raise ValueError("Invalid config")
config = TestConfig(value="test")
artifact_storage = ArtifactStorage(artifact_path=tmp_path)
resource_provider = ResourceProvider(artifact_storage=artifact_storage)
task = TestTask(config=config, resource_provider=resource_provider)
assert task._config.value == "test"
invalid_config = TestConfig(value="invalid")
with pytest.raises(ValueError, match="Invalid config"):
TestTask(config=invalid_config, resource_provider=resource_provider)
def test_configurable_task_resource_validation(tmp_path) -> None:
class TestConfig(ConfigBase):
value: str
class TestTask(ConfigurableTask[TestConfig]):
@classmethod
def get_config_type(cls) -> type[TestConfig]:
return TestConfig
def _validate(self) -> None:
pass
def _initialize(self) -> None:
pass
config = TestConfig(value="test")
artifact_storage = ArtifactStorage(artifact_path=tmp_path)
mock_model_registry = Mock(spec=ModelRegistry)
resource_provider = ResourceProvider(artifact_storage=artifact_storage, model_registry=mock_model_registry)
task = TestTask(config=config, resource_provider=resource_provider)
assert task._resource_provider == resource_provider
def test_configurable_task_resource_provider_is_none() -> None:
class TestConfig(ConfigBase):
value: str
class TestTask(ConfigurableTask[TestConfig]):
def _validate(self) -> None:
pass
def _initialize(self) -> None:
pass
config = TestConfig(value="test")
task = TestTask(config=config, resource_provider=None)
assert task._resource_provider is None