fix: bridge model.generate() to agenerate() for custom columns in async engine

[andreatgretel]

📋 Summary

Custom column generators that call model.generate() inside their function body fail under the
async engine because the sync HTTP client is unavailable. This adds a transparent proxy that
bridges to model.agenerate() via run_coroutine_threadsafe, so user code works unchanged
in both engines.

🔗 Related Issue

N/A - discovered via NVIDIA-NeMo/Anonymizer#119 where a model_compat.model_generate() workaround
was needed; this PR moves the fix into DD so all consumers get it for free.

🔄 Changes

• Add _AsyncBridgedModelFacade proxy class that intercepts the sync-client RuntimeError and
  schedules agenerate() on the engine's persistent event loop
  (custom.py#L25-L72)
• Wrap facades in _build_models_dict() so the bridge is transparent to user code
• Include deadlock guard: clear error if called from the event loop instead of hanging
• Match only the exact HttpModelClient sync-mode error (not substring matching)

🧪 Testing

• make test passes (197 column generator tests)
• Unit tests added: 4 tests covering sync passthrough, async bridge, error propagation, deadlock guard
• E2E tests added/updated — N/A, no end-to-end test infrastructure for async engine mode

✅ Checklist

• Follows commit message conventions
• Commits are signed off (DCO)
• Architecture docs updated — N/A, internal proxy class

[github-actions]

Code Review: PR #545 — fix: bridge model.generate() to agenerate() for custom columns in async engine

Summary

This PR adds an _AsyncBridgedModelFacade proxy class that transparently bridges synchronous model.generate() calls to model.agenerate() when a custom column generator runs inside the async engine. Under the async engine, sync custom columns execute via asyncio.to_thread, where the sync HTTP client is unavailable. The proxy intercepts the specific RuntimeError from the async-mode HttpModelClient and schedules agenerate() on the engine's persistent event loop via run_coroutine_threadsafe. A deadlock guard prevents misuse from the event loop thread itself.

Files changed: 2 (1 source, 1 test) — +160 / -3 lines.

Findings

Correctness

1. Exact error string matching is fragile (Medium)
   custom.py:35 — The bridge triggers only when str(exc) == "Sync methods are not available on an async-mode HttpModelClient.". If the error message is ever reworded upstream (even a punctuation change), the bridge silently stops working and users get an opaque RuntimeError. Consider importing and matching on a specific exception subclass from the HTTP client, or at minimum matching a stable substring/prefix. If the upstream error message is under this project's control, adding a comment cross-referencing the source would help.

2. kwargs not forwarded in test_bridges_to_agenerate_on_sync_client_error (Low)
   test_custom.py:533 — The test calls proxy.generate("hello", parser=str) and asserts result == ("async_result", ["hello"]), but the fake_agenerate mock (*args, **kwargs) silently drops parser=str from the assertion. The test doesn't verify that keyword arguments survive the bridge. Consider asserting on kwargs too, e.g., returning (args, kwargs) and checking both.

3. Timeout not tested (Low)
   There is no test for the _SYNC_BRIDGE_TIMEOUT path. If agenerate() hangs, the future.result(timeout=...) will raise concurrent.futures.TimeoutError, which will propagate as-is rather than as a project-canonical error. This matches the existing pattern in base.py:45-49 which wraps the timeout, but the facade does not. Consider whether a TimeoutError leaking from this path is acceptable or should be wrapped.

Design

4. Proxy is always applied, even in sync engine mode (Info)
   custom.py:337 — _build_models_dict() unconditionally wraps every facade in _AsyncBridgedModelFacade. In sync engine mode the proxy's generate() calls facade.generate() which succeeds immediately, so the overhead is one extra try/except per call. This is a reasonable trade-off for simplicity, but worth noting.

5. __slots__ + __getattr__ proxy pattern is clean (Positive)
   The use of __slots__ with object.__setattr__ / object.__getattribute__ for the internal _facade attribute, combined with __getattr__ forwarding, is a solid proxy pattern that avoids attribute shadowing issues. The __repr__ is helpful for debugging.

6. Deferred import of ensure_async_engine_loop inside generate() (Info)
   custom.py:62 — The import is deferred to avoid a circular import or heavy import at module load time. This is consistent with the project's lazy-import convention. Fine as-is.

Testing

7. Good coverage of key scenarios (Positive)
   Four tests cover: sync passthrough, async bridge, error propagation for non-matching errors (both different exception type and different RuntimeError message), and the deadlock guard. The deadlock guard test (test_deadlock_guard_on_event_loop) correctly uses asyncio.run() to simulate calling from the event loop.

8. Test creates a real event loop + thread (Info)
   test_custom.py:523-537 — The async bridge test spins up a real event loop on a background thread, which is the correct approach for testing run_coroutine_threadsafe. The cleanup in finally is proper (stop + join with timeout).

Style / Conventions

9. Import of _SYNC_BRIDGE_TIMEOUT from base (Info)
   The constant is reused from base.py rather than duplicated, which is good. The leading underscore signals it's internal, consistent with its use in base.py.

10. from __future__ import annotations present in source file (Positive)
    Follows the project's structural invariant.

Verdict

Approve with minor suggestions. The implementation is well-structured, follows project conventions, and includes solid test coverage. The proxy pattern is clean, the deadlock guard is a thoughtful addition, and the tests cover the important scenarios.

The main concern is the exact-string error matching (Finding #1), which creates a fragile coupling to an upstream error message. This is acceptable for a short-term fix but should be tracked for improvement — ideally by introducing a typed exception in the HTTP client layer. The missing kwargs assertion in the bridge test (Finding #2) is minor but easy to fix.

No blocking issues found.

[greptile-apps]

Greptile Summary

Introduces _AsyncBridgedModelFacade, a thin proxy that intercepts the HttpModelClient sync-mode RuntimeError in model.generate() and bridges it to model.agenerate() via run_coroutine_threadsafe on the engine's singleton event loop, making sync custom columns work transparently under the async engine. The proxy is always applied in _build_models_dict() and is a no-op in sync mode (the first facade.generate() call succeeds and returns immediately). The deadlock guard, exact error matching, and lazy import of ensure_async_engine_loop are all handled correctly.

Confidence Score: 5/5

• This PR is safe to merge — the proxy is logically correct, transparent in sync mode, and all four new tests pass the key paths including the deadlock guard.
• No P0 or P1 findings. The bridge logic is sound: exact error matching prevents false positives, the running-loop check prevents deadlocks, the singleton engine loop is reused correctly, and the lazy import avoids circular dependencies. Tests cover all meaningful branches.
• No files require special attention.

Important Files Changed

Filename	Overview
packages/data-designer-engine/src/data_designer/engine/column_generators/generators/custom.py	Adds _AsyncBridgedModelFacade proxy and updates _build_models_dict to wrap all facades; logic is correct — sync pass-through, async bridge via run_coroutine_threadsafe, and deadlock guard all behave as intended.
packages/data-designer-engine/tests/engine/column_generators/generators/test_custom.py	Adds four targeted tests (sync pass-through, async bridge, non-client-error propagation, deadlock guard); coverage is complete for the new proxy class and the patch target is correct.

Sequence Diagram

sequenceDiagram
    participant UC as User Custom Column
    participant P as _AsyncBridgedModelFacade
    participant F as ModelFacade
    participant L as Engine Event Loop<br/>(ensure_async_engine_loop)

    Note over UC,L: Sync Engine Path
    UC->>P: model.generate(prompt)
    P->>F: facade.generate(prompt)
    F-->>P: (result, metadata)
    P-->>UC: (result, metadata)

    Note over UC,L: Async Engine Path (worker thread via asyncio.to_thread)
    UC->>P: model.generate(prompt)
    P->>F: facade.generate(prompt)
    F-->>P: RuntimeError("Sync methods are not available...")
    P->>P: asyncio.get_running_loop() → RuntimeError (no loop in worker thread)
    P->>L: ensure_async_engine_loop()
    L-->>P: persistent event loop
    P->>L: run_coroutine_threadsafe(facade.agenerate(prompt))
    L->>F: await facade.agenerate(prompt)
    F-->>L: (result, metadata)
    L-->>P: future resolved
    P-->>UC: (result, metadata)

    Note over UC,L: Deadlock Guard (called from event loop)
    UC->>P: model.generate(prompt)
    P->>F: facade.generate(prompt)
    F-->>P: RuntimeError("Sync methods are not available...")
    P->>P: asyncio.get_running_loop() → loop (running!)
    P-->>UC: RuntimeError("Use 'await model.agenerate()'")

Loading

_{Reviews (1): Last reviewed commit: "feat: bridge model.generate() to agenera..." | Re-trigger Greptile}