chore: lazy 3rd party imports (#222)

Author: nabinchha

.gitignore

@@ -93,3 +93,6 @@ docs/notebook_source/*.csv
 docs/**/artifacts/
 
 tests_e2e/uv.lock
+
+# Performance profiling
+perf_*.txt

AGENTS.md

@@ -158,12 +158,13 @@ Type annotations are REQUIRED for all code in this project. This is strictly enf
 ### Import Style
 
 - **ALWAYS** use absolute imports, never relative imports
-- Place imports at module level, not inside functions
+- Place imports at module level, not inside functions (exception: it is unavoidable for performance reasons)
 - Import sorting is handled by `ruff`'s `isort` - imports should be grouped and sorted:
   1. Standard library imports
-  2. Third-party imports
+  2. Third-party imports (use `lazy_heavy_imports` for heavy libraries)
   3. First-party imports (`data_designer`)
 - Use standard import conventions (enforced by `ICN`)
+- See [Lazy Loading and TYPE_CHECKING](#lazy-loading-and-type_checking) section for optimization guidelines
 
   ```python
   # Good
@@ -184,6 +185,146 @@ Type annotations are REQUIRED for all code in this project. This is strictly enf
       path = Path(filename)
   ```
 
+### Lazy Loading and TYPE_CHECKING
+
+This project uses lazy loading for heavy third-party dependencies to optimize import performance.
+
+#### When to Use Lazy Loading
+
+**Heavy third-party libraries** (>100ms import cost) should be lazy-loaded via `lazy_heavy_imports.py`:
+
+```python
+# ❌ Don't import directly
+import pandas as pd
+import numpy as np
+
+# ✅ Use lazy loading with IDE support
+from typing import TYPE_CHECKING
+from data_designer.lazy_heavy_imports import pd, np
+
+if TYPE_CHECKING:
+    import pandas as pd  # For IDE autocomplete and type hints
+    import numpy as np
+```
+
+This pattern provides:
+- Runtime lazy loading (fast startup)
+- Full IDE support (autocomplete, type hints)
+- Type checker validation
+
+**See [lazy_heavy_imports.py](src/data_designer/lazy_heavy_imports.py) for the current list of lazy-loaded libraries.**
+
+#### Adding New Heavy Dependencies
+
+If you add a new dependency with significant import cost (>100ms):
+
+1. **Add to `lazy_heavy_imports.py`:**
+   ```python
+   _LAZY_IMPORTS = {
+       # ... existing entries ...
+       "your_lib": "your_library_name",
+   }
+   ```
+
+2. **Update imports across codebase:**
+   ```python
+   from typing import TYPE_CHECKING
+   from data_designer.lazy_heavy_imports import your_lib
+
+   if TYPE_CHECKING:
+       import your_library_name as your_lib  # For IDE support
+   ```
+
+3. **Verify with performance test:**
+   ```bash
+   make perf-import CLEAN=1
+   ```
+
+#### Using TYPE_CHECKING Blocks
+
+`TYPE_CHECKING` blocks defer imports that are only needed for type hints, preventing circular dependencies and reducing import time.
+
+**For internal data_designer imports:**
+
+```python
+from __future__ import annotations  # Always include at top
+
+from typing import TYPE_CHECKING
+
+# Runtime imports
+from pathlib import Path
+from data_designer.config.base import ConfigBase
+
+if TYPE_CHECKING:
+    # Type-only imports - only visible to type checkers
+    from data_designer.engine.models.facade import ModelFacade
+
+def get_model(model: ModelFacade) -> str:
+    return model.name
+```
+
+**For lazy-loaded libraries (see pattern in "When to Use Lazy Loading" above):**
+- Import from `lazy_heavy_imports` for runtime
+- Add full import in `TYPE_CHECKING` block for IDE support
+
+**Rules for TYPE_CHECKING:**
+
+✅ **DO put in TYPE_CHECKING:**
+- Internal `data_designer` imports used **only** in type hints
+- Imports that would cause circular dependencies
+- **Full imports of lazy-loaded libraries for IDE support** (e.g., `import pandas as pd` in addition to runtime `from data_designer.lazy_heavy_imports import pd`)
+
+❌ **DON'T put in TYPE_CHECKING:**
+- **Standard library imports** (`Path`, `Any`, `Callable`, `Literal`, `TypeAlias`, etc.)
+- **Pydantic model types** used in field definitions (needed at runtime for validation)
+- **Types used in discriminated unions** (Pydantic needs them at runtime)
+- **Any import used at runtime** (instantiation, method calls, base classes, etc.)
+
+**Examples:**
+
+```python
+# ✅ CORRECT - Lazy-loaded library with IDE support
+from typing import TYPE_CHECKING
+from data_designer.lazy_heavy_imports import pd
+
+if TYPE_CHECKING:
+    import pandas as pd  # IDE gets full type hints
+
+def load_data(path: str) -> pd.DataFrame:  # IDE understands pd.DataFrame
+    return pd.read_csv(path)
+
+# ✅ CORRECT - Standard library NOT in TYPE_CHECKING
+from pathlib import Path
+from typing import Any
+
+def process_file(path: Path) -> Any:
+    return path.read_text()
+
+# ✅ CORRECT - Internal type-only import
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from data_designer.engine.models.facade import ModelFacade
+
+def get_model(model: ModelFacade) -> str:  # Only used in type hint
+    return model.name
+
+# ❌ INCORRECT - Pydantic field type in TYPE_CHECKING
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from data_designer.config.models import ModelConfig  # Wrong!
+
+class MyConfig(BaseModel):
+    model: ModelConfig  # Pydantic needs this at runtime!
+
+# ✅ CORRECT - Pydantic field type at runtime
+from data_designer.config.models import ModelConfig
+
+class MyConfig(BaseModel):
+    model: ModelConfig
+```
+
 ### Naming Conventions (PEP 8)
 
 Follow PEP 8 naming conventions:

Makefile

@@ -45,14 +45,25 @@ help:
 	@echo "  check-license-headers     - Check if all files have license headers"
 	@echo "  update-license-headers    - Add license headers to all files"
 	@echo ""
+	@echo "⚡ Performance:"
+	@echo "  perf-import               - Profile import time and show summary"
+	@echo "  perf-import CLEAN=1       - Clean cache, then profile import time"
+	@echo "  perf-import NOFILE=1      - Profile without writing to file (for CI)"
+	@echo ""
 	@echo "═════════════════════════════════════════════════════════════"
 	@echo "💡 Tip: Run 'make <command>' to execute any command above"
 	@echo ""
 
-clean:
-	@echo "🧹 Cleaning up coverage reports and cache files..."
-	rm -rf htmlcov .coverage .pytest_cache
+clean-pycache:
+	@echo "🧹 Cleaning up Python cache files..."
 	find . -type d -name __pycache__ -exec rm -rf {} + 2>/dev/null || true
+	find . -type f -name "*.pyc" -delete 2>/dev/null || true
+	@echo "✅ Cache cleaned!"
+
+clean: clean-pycache
+	@echo "🧹 Cleaning up coverage reports and test cache..."
+	rm -rf htmlcov .coverage .pytest_cache
+	@echo "✅ Cleaned!"
 
 coverage:
 	@echo "📊 Running tests with coverage analysis..."
@@ -168,4 +179,34 @@ install-dev-notebooks:
 	$(call install-pre-commit-hooks)
 	@echo "✅ Dev + notebooks installation complete!"
 
-.PHONY: clean coverage format format-check lint lint-fix test test-e2e test-run-tutorials test-run-recipes test-run-all-examples check-license-headers update-license-headers check-all check-all-fix install install-dev install-dev-notebooks generate-colab-notebooks
+perf-import:
+ifdef CLEAN
+	@$(MAKE) clean-pycache
+endif
+	@echo "⚡ Profiling import time for data_designer.essentials..."
+ifdef NOFILE
+	@PERF_OUTPUT=$$(uv run python -X importtime -c "import data_designer.essentials" 2>&1); \
+	echo "$$PERF_OUTPUT"; \
+	echo ""; \
+	echo "Summary:"; \
+	echo "$$PERF_OUTPUT" | tail -1 | awk '{printf "  Total: %.3fs\n", $$5/1000000}'; \
+	echo ""; \
+	echo "💡 Top 10 slowest imports:"; \
+	printf "%-12s %-12s %s\n" "Self (s)" "Cumulative (s)" "Module"; \
+	printf "%-12s %-12s %s\n" "--------" "--------------" "------"; \
+	echo "$$PERF_OUTPUT" | grep "import time:" | sort -rn -k5 | head -10 | awk '{printf "%-12.3f %-12.3f %s", $$3/1000000, $$5/1000000, $$7; for(i=8;i<=NF;i++) printf " %s", $$i; printf "\n"}'
+else
+	@PERF_FILE="perf_import_$$(date +%Y%m%d_%H%M%S).txt"; \
+	uv run python -X importtime -c "import data_designer.essentials" > "$$PERF_FILE" 2>&1; \
+	echo "📊 Import profile saved to $$PERF_FILE"; \
+	echo ""; \
+	echo "Summary:"; \
+	tail -1 "$$PERF_FILE" | awk '{printf "  Total: %.3fs\n", $$5/1000000}'; \
+	echo ""; \
+	echo "💡 Top 10 slowest imports:"; \
+	printf "%-12s %-12s %s\n" "Self (s)" "Cumulative (s)" "Module"; \
+	printf "%-12s %-12s %s\n" "--------" "--------------" "------"; \
+	grep "import time:" "$$PERF_FILE" | sort -rn -k5 | head -10 | awk '{printf "%-12.3f %-12.3f %s", $$3/1000000, $$5/1000000, $$7; for(i=8;i<=NF;i++) printf " %s", $$i; printf "\n"}'
+endif
+
+.PHONY: clean clean-pycache coverage format format-check lint lint-fix test test-e2e test-run-tutorials test-run-recipes test-run-all-examples check-license-headers update-license-headers check-all check-all-fix install install-dev install-dev-notebooks generate-colab-notebooks perf-import

src/data_designer/__init__.py

@@ -1,6 +1,8 @@
 # SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: Apache-2.0
 
+from __future__ import annotations
+
 try:
     from data_designer._version import __version__
 except ImportError:

src/data_designer/cli/__init__.py

@@ -1,6 +1,8 @@
 # SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: Apache-2.0
 
+from __future__ import annotations
+
 from data_designer.cli.main import app, main
 
 __all__ = ["app", "main"]

src/data_designer/cli/commands/download.py

@@ -1,6 +1,8 @@
 # SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: Apache-2.0
 
+from __future__ import annotations
+
 import typer
 
 from data_designer.cli.controllers.download_controller import DownloadController

src/data_designer/cli/commands/list.py

@@ -1,6 +1,8 @@
 # SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: Apache-2.0
 
+from __future__ import annotations
+
 from rich.table import Table
 
 from data_designer.cli.repositories.model_repository import ModelRepository

src/data_designer/cli/commands/models.py

@@ -1,6 +1,8 @@
 # SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: Apache-2.0
 
+from __future__ import annotations
+
 from data_designer.cli.controllers.model_controller import ModelController
 from data_designer.config.utils.constants import DATA_DESIGNER_HOME
 

src/data_designer/cli/commands/providers.py

@@ -1,6 +1,8 @@
 # SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: Apache-2.0
 
+from __future__ import annotations
+
 from data_designer.cli.controllers.provider_controller import ProviderController
 from data_designer.config.utils.constants import DATA_DESIGNER_HOME
 

src/data_designer/cli/commands/reset.py

@@ -1,6 +1,8 @@
 # SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: Apache-2.0
 
+from __future__ import annotations
+
 import typer
 
 from data_designer.cli.repositories.model_repository import ModelRepository