fix: break health check import cycle

[andreatnvidia]

📋 Summary

Restore provider health checks without exposing engine-internal retry taxonomy to the external CI script. DataDesigner.check_models now owns retry classification and backoff behind its public interface while preserving one-shot behavior by default.

🔗 Related Issue

N/A - regression exposed by #750 and observed in Health Checks #48

🔄 Changes

• Add validated, opt-in retry settings to DataDesigner.check_models
• Retry only canonical transient model errors and readiness timeouts with linear backoff
• Keep scripts/health_checks.py on the public interface with three attempts and 5s/10s delays
• Add coverage for retryable, non-retryable, exhausted, invalid, and fresh-process entrypoint cases

🧪 Testing

• .venv/bin/ruff check --fix .
• .venv/bin/ruff format .
• .venv/bin/pytest packages/data-designer/tests -q - 988 passed, 1 skipped
• No-secret health-check smoke test - 0 failed, 12 skipped
• E2E tests added/updated (N/A - import and retry boundary regression)

✅ Checklist

• Follows commit message conventions
• Commits are signed off (DCO)
• Architecture docs updated (N/A - no architecture change)

[greptile-apps]

Greptile Summary

This PR fixes an import cycle introduced by #750, where scripts/health_checks.py was directly importing engine-internal retry taxonomy (RETRYABLE_MODEL_ERRORS) from data_designer.engine.models.errors. The fix moves retry classification and backoff into DataDesigner.check_models as opt-in parameters, keeping the CI script on the stable public interface.

• run_readiness_check drops the client_concurrency_mode parameter and unconditionally uses the async path; the ensure_async_engine_loop import is promoted from a lazy conditional to a top-level import.
• DataDesigner.check_models gains max_attempts (default 1, one-shot) and retry_backoff_seconds (default 0) with linear backoff; retries are scoped to RETRYABLE_MODEL_ERRORS + (TimeoutError,) exactly matching the old CI script's behavior.
• scripts/health_checks.py is simplified to pass max_attempts=3, retry_backoff_seconds=5 through the new public interface, eliminating its direct engine import; a fresh-process import test is added to guard against the cycle recurring.

Confidence Score: 5/5

Safe to merge. The change is a targeted refactor with no behavioural regressions: retry semantics are identical to the old CI script, the public one-shot default is preserved, and all edge cases are covered by new tests.

The retry logic, backoff formula, and error taxonomy are all correct and verified by the new parametrised tests. The async-only simplification in run_readiness_check is consistent with how both call sites (dataset builder and check_models) already invoked it. No issues were found across any of the seven changed files.

No files require special attention.

Important Files Changed

Filename	Overview
packages/data-designer/src/data_designer/interface/data_designer.py	Adds opt-in retry to check_models with validated max_attempts / retry_backoff_seconds and linear backoff; retry taxonomy (RETRYABLE_MODEL_ERRORS + TimeoutError) moved from the CI script into this public interface.
packages/data-designer-engine/src/data_designer/engine/readiness.py	Removes the now-unused client_concurrency_mode parameter and always uses the async path; ensure_async_engine_loop promoted to a top-level import. Logic is correct and cleaner.
scripts/health_checks.py	Simplified: no longer imports engine-internal RETRYABLE_MODEL_ERRORS or manages its own retry loop; delegates retry to DataDesigner.check_models. Correct and minimal.
packages/data-designer-engine/tests/engine/test_readiness.py	Replaces inline per-test async mocking with a consolidated autouse fixture; all assertions updated to arun_health_check. Coverage is complete and logically sound.
packages/data-designer/tests/interface/test_data_designer.py	New tests cover retryable errors, non-retryable early propagation, exhausted-attempt re-raise, sleep call sequencing, and input validation; fixtures refactored to shared scope.
packages/data-designer/tests/test_lazy_imports.py	New test_health_checks_entrypoint_imports_in_fresh_process test verifies the import cycle is resolved by running the script in a subprocess with all API keys removed and asserting exit code 0.
packages/data-designer-engine/src/data_designer/engine/dataset_builders/dataset_builder.py	Two call sites of run_readiness_check drop the now-removed client_concurrency_mode kwarg and the associated import; straightforward cleanup.

Sequence Diagram

%%{init: {'theme': 'neutral'}}%%
sequenceDiagram
    participant Script as health_checks.py
    participant DD as DataDesigner.check_models
    participant RR as run_readiness_check
    participant MR as ModelRegistry.arun_health_check
    participant MCP as MCPRegistry.run_health_check

    Script->>DD: "check_models(config, max_attempts=3, retry_backoff_seconds=5)"
    loop attempt 1..max_attempts
        DD->>RR: run_readiness_check(columns, resource_provider)
        RR->>MR: arun_health_check(aliases) [via run_coroutine_threadsafe]
        MR-->>RR: result / error
        RR->>MCP: run_health_check(tool_aliases)
        MCP-->>RR: ok / error
        RR-->>DD: ok / retryable error / non-retryable error
        alt non-retryable error
            DD-->>Script: raise immediately
        else retryable error AND attempts exhausted
            DD-->>Script: raise
        else retryable error AND attempts remain
            DD->>DD: "time.sleep(attempt * backoff)"
        else success
            DD-->>Script: return
        end
    end

Loading

%%{init: {'theme': 'base', 'themeVariables': {"darkMode": true, "background": "#0d1117", "primaryColor": "#21262d", "primaryTextColor": "#e6edf3", "primaryBorderColor": "#8b949e", "lineColor": "#8b949e", "textColor": "#e6edf3", "edgeLabelBackground": "#161b22", "actorBkg": "#21262d", "actorBorder": "#8b949e", "actorTextColor": "#e6edf3", "actorLineColor": "#8b949e", "signalColor": "#8b949e", "signalTextColor": "#e6edf3", "noteBkgColor": "#373320", "noteBorderColor": "#d4a72c", "noteTextColor": "#f0e6c0", "labelBoxBkgColor": "#21262d", "labelBoxBorderColor": "#8b949e", "labelTextColor": "#e6edf3", "loopTextColor": "#e6edf3", "activationBkgColor": "#30363d", "activationBorderColor": "#8b949e"}}}%%
sequenceDiagram
    participant Script as health_checks.py
    participant DD as DataDesigner.check_models
    participant RR as run_readiness_check
    participant MR as ModelRegistry.arun_health_check
    participant MCP as MCPRegistry.run_health_check

    Script->>DD: "check_models(config, max_attempts=3, retry_backoff_seconds=5)"
    loop attempt 1..max_attempts
        DD->>RR: run_readiness_check(columns, resource_provider)
        RR->>MR: arun_health_check(aliases) [via run_coroutine_threadsafe]
        MR-->>RR: result / error
        RR->>MCP: run_health_check(tool_aliases)
        MCP-->>RR: ok / error
        RR-->>DD: ok / retryable error / non-retryable error
        alt non-retryable error
            DD-->>Script: raise immediately
        else retryable error AND attempts exhausted
            DD-->>Script: raise
        else retryable error AND attempts remain
            DD->>DD: "time.sleep(attempt * backoff)"
        else success
            DD-->>Script: return
        end
    end

Loading

_{Reviews (4): Last reviewed commit: "Merge branch 'main' into andreatgretel/f..." | Re-trigger Greptile}

[github-actions]

Code Review: PR #781 — fix: break health check import cycle

Author: andreatgretel (Andre Manoel) · Base: main · Files changed: 2 · +34 / -1

Summary

This PR fixes a circular-import regression that disabled provider health checks. The cycle arose when data_designer.engine.models.errors was imported before the model-client factory: the clients package (engine/models/clients/__init__.py) eagerly imported create_model_client from factory.py, and factory.py in turn imports back into the models/clients subtree (including retry, errors, and the model_request_executor), so importing the error constants first triggered a partially-initialized module.

The fix removes the eager from ...factory import create_model_client and re-exports it lazily via a module-level __getattr__, importing factory only on first attribute access. A TYPE_CHECKING import preserves static type-checker / IDE visibility, and create_model_client remains in __all__. A fresh-process regression test pins the previously-failing import order.

This is a clean, minimal, well-targeted fix.

Findings

Correctness — looks right

• Lazy resolution matches the established pattern. The new __getattr__/__dir__ mirrors the existing lazy-export implementations in interface/__init__.py:43 and config/__init__.py:247 — same importlib.import_module + globals()[name] = attr caching so subsequent accesses bypass __getattr__, same AttributeError fallback for unknown names. Consistency with surrounding code is good.
• __all__ membership preserved. create_model_client is kept in __all__ (line 56), so from data_designer.engine.models.clients import * still resolves it (Python invokes __getattr__ for __all__ names not found in the namespace), and __dir__() correctly advertises it for tab-completion. No public-API surface change — the docstring claim "remains available through the same package export" holds.
• TYPE_CHECKING guard is correct. The # noqa: F401 is appropriate since the guarded import exists only for type checkers. from __future__ import annotations is already present, so this composes correctly.
• Cycle is genuinely broken, not merely deferred to a different trigger: factory.py is no longer touched at clients-package import time, so importing engine.models.errors (or its RETRYABLE_MODEL_ERRORS, a 4-tuple — matching the test's expected "4") no longer drags in factory transitively.

Style / conventions — minor, non-blocking

• Single-name __getattr__ vs. the _LAZY_IMPORTS dict. The two sibling modules (interface, config) drive lazy exports from a _LAZY_IMPORTS: dict[str, tuple[str, str]] table. This PR hardcodes the one name with an if name == "create_model_client": branch. For a single export that's reasonable and arguably clearer, but if more lazy re-exports are added here later, converging on the dict pattern would keep the three modules uniform. Not worth changing now.
• __dir__ returns __all__ directly. Consistent with the siblings; fine. (Returning the bare list rather than a copy is harmless here since callers don't mutate it, and the siblings do the same.)

Tests

• The new test_model_errors_import_before_client_factory_in_fresh_process reproduces the exact failing order in a fresh subprocess via the existing _run_python_snippet helper — the right approach, since import-cycle bugs are order- and process-state-dependent and would be masked inside an already-warmed interpreter.
• It asserts both that the errors constant loads (len(RETRYABLE_MODEL_ERRORS) == 4) and that create_model_client resolves to the real callable (.__name__), so it guards against a regression where the lazy hook returns the wrong object or silently shadows the name. Good coverage for the size of the change.
• Per the PR body, full engine (2178) and interface (980) suites pass plus a no-secret health-check smoke test; I was unable to re-run them or ruff in this CI shell (package not importable / ruff not on PATH here), so I'm relying on the reported runs for those.

Security / performance

• No security implications.
• Performance: marginally improves cold-import cost — factory.py and its dependency fan-out are no longer imported when only the error constants or other client types are needed. The one-time import_module on first create_model_client access is negligible and cached thereafter.

Suggestions (optional, none blocking)

1. Consider a one-line comment above __getattr__ noting why the factory is lazy (the cycle), so a future refactor doesn't "tidy" it back into an eager import and silently reintroduce the regression. The test guards against it, but an inline breadcrumb at the edit site helps.
2. If additional lazy re-exports land in this module, migrate to the _LAZY_IMPORTS dict form used by the sibling packages for uniformity.

Structural Impact (graphify, 2.4s)

Risk: HIGH (2 import direction violation(s))

• 2 Python files, 3 AST entities, 1/82 clusters

Import Direction Violations (2)

Legal direction: interface -> engine -> config

• Lazily import config module exports when accessed. (config) --rationale_for--> __getattr__() (engine)
• Return list of available exports for tab-completio (config) --rationale_for--> __dir__() (engine)

High-Connectivity Changes

• __getattr__() (6 deps) in packages/data-designer-engine/src/data_designer/engine/models/clients/__init__.py
• __dir__() (5 deps) in packages/data-designer-engine/src/data_designer/engine/models/clients/__init__.py

Cross-Package Dependencies

• Lazily import interface exports when accessed. (interface) --rationale_for--> __getattr__() (engine)
• Return available exports for tab-completion. (interface) --rationale_for--> __dir__() (engine)
• Lazily import config module exports when accessed. (config) --rationale_for--> __getattr__() (engine)
• Return list of available exports for tab-completio (config) --rationale_for--> __dir__() (engine)

Reviewer note on structural impact: The HIGH risk score and the two flagged "import direction violations" appear to be false positives from naming-based similarity matching, not real architectural problems. The tool is associating this __getattr__/__dir__ pair with the identically-named lazy-import functions in the config and interface packages (the "rationale_for" edges all point between same-named __getattr__/__dir__ definitions). The actual diff introduces no new cross-package imports: the only new import is data_designer.engine.models.clients.factory (engine → engine, deferred), which respects the interface → engine → config direction. I reviewed the import graph manually and found no reverse imports. The genuine signal worth heeding is the blast radius of create_model_client (a high-connectivity, widely-depended-upon export) — and on that front, keeping it in __all__ plus the fresh-process regression test adequately protect backward compatibility.

Verdict

Approve (non-blocking suggestions only). Small, surgical fix that resolves a real regression, follows the codebase's existing lazy-export idiom, preserves the public API surface, and adds a targeted regression test that reproduces the exact failure mode. No correctness, security, or backward-compatibility concerns. The structural-impact HIGH rating does not reflect a real import-direction violation (see reviewer note above). I did not re-run the test suites or linters in this environment; the verdict trusts the passing runs reported in the PR description.

[nabinchha]

Thanks for putting this together, @andreatgretel!

Summary

This restores the data_designer.engine.models.clients.create_model_client package export without eagerly importing the factory, which avoids the import cycle triggered when model retry errors are imported first. I agree this unblocks the health-check regression, but I think the current fix treats the symptom rather than the boundary issue that caused the cycle.

Findings

Warnings — Worth addressing

Design issues, missing error handling, test gaps, or violations of project standards that could cause problems later.

packages/data-designer-engine/src/data_designer/engine/models/clients/__init__.py:62 — Avoid using a lazy package export to compensate for a leaked engine boundary

• What: This PR fixes the cycle by making create_model_client lazy at the clients package boundary. The immediate trigger, though, is that the provider health-check script added in the previous PR imports data_designer.engine.models.errors.RETRYABLE_MODEL_ERRORS directly so it can classify retries outside the public health-check/readiness path.
• Why: That means a top-level CI script now knows about engine-internal model retry taxonomy. The lazy re-export is a reasonable facade pattern in isolation, but here it lets the cross-boundary dependency remain in place and makes import order part of the architecture. Future health-check callers could repeat the same pattern instead of relying on a single readiness contract.
• Suggestion: Could we move the retry policy behind the health-check boundary instead? For example, DataDesigner.check_models(...), engine.readiness.run_readiness_check(...), or a small public helper could own retry classification and backoff. Then scripts/health_checks.py can keep calling the public interface without importing engine.models.errors, and we may not need to change the clients package export at all.

What Looks Good

• The fix keeps the public from data_designer.engine.models.clients import create_model_client surface intact while deferring only the cycle-forming factory import.
• The implementation follows the same __getattr__/cache pattern already used by data_designer.interface and data_designer.config.
• The new regression test exercises the import order in a fresh Python process, which is the right shape for circular-import coverage.

Verdict

Needs changes — I’d prefer we address the health-check retry boundary directly instead of adding a lazy export primarily to make the import cycle survivable.

---

This review was generated by an AI assistant.

[github-actions]

Fern preview: https://nvidia-preview-pr-781.docs.buildwithfern.com/nemo/datadesigner

Fern previews include the docs-website version archive with PR changes synced into latest. Notebook tutorials are rendered without execution outputs in previews.

[andreatnvidia]

Thanks for the review, @nabinchha. I pushed e3c344ae to:

• Move retry classification and backoff into DataDesigner.check_models, preserving one-shot behavior by default.
• Remove the engine-internal error import from scripts/health_checks.py.
• Revert the lazy create_model_client export.
• Add retry and fresh-process entrypoint coverage.

CI is rerunning now.

[nabinchha]

oh did we miss this in the PR to drop SYNC? we shouldn't need this distinction right?

[andreatnvidia]

@nabinchha Yep, this was another leftover from #767. I pushed 9cac4619 to make readiness async-only: run_readiness_check now always uses arun_health_check, callers no longer pass a concurrency mode, and the obsolete sync-readiness test is removed.

I also did a repo-wide pass and found a few other sync-engine leftovers outside the readiness path. They do not affect this fix, so we’ll address them in a separate cleanup PR.

Local validation: engine 2170 passed; interface 988 passed, 1 skipped; Ruff clean.