data-designer-engine 0.9.6__tar.gz → 0.9.7__tar.gz

{data_designer_engine-0.9.6 → data_designer_engine-0.9.7}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.5
 Name: data-designer-engine
-Version: 0.9.6
+Version: 0.9.7
 Summary: Generation engine for DataDesigner synthetic data generation
 License-Expression: Apache-2.0
 Classifier: Development Status :: 4 - Beta
@@ -16,7 +16,7 @@ Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
 Requires-Python: >=3.10
 Requires-Dist: anyascii<1,>=0.3.3
 Requires-Dist: asciichartpy<2,>=1.5.25
-Requires-Dist: data-designer-config==0.9.6
+Requires-Dist: data-designer-config==0.9.7
 Requires-Dist: duckdb<2,>=1.5.0
 Requires-Dist: faker<21,>=20.1.0
 Requires-Dist: fsspec<2027,>=2026.6.0

{data_designer_engine-0.9.6 → data_designer_engine-0.9.7}/src/data_designer/engine/dataset_builders/utils/async_concurrency.py

@@ -19,12 +19,19 @@ Startup handshake:
     ``_loop.is_running() == False`` before the first loop has entered
     ``run_forever()``, which would create a duplicate loop. On timeout,
     globals are reset and the orphaned loop is cleaned up before raising.
+
+Fork safety:
+    Only the forking thread survives ``os.fork()``, but the copied loop
+    object still reports ``is_running() == True``. An ``after_in_child``
+    hook drops the inherited loop, thread, and lock so the child's first
+    call starts a fresh loop instead of waiting on one nothing drives.
 """
 
 from __future__ import annotations
 
 import asyncio
 import logging
+import os
 import threading
 
 logger = logging.getLogger(__name__)
@@ -42,6 +49,19 @@ def _run_loop(loop: asyncio.AbstractEventLoop, ready: threading.Event) -> None:
     loop.run_forever()
 
 
+def _reset_after_fork() -> None:
+    # Drop, never close(), the inherited loop: the child shares the parent's epoll
+    # instance, and close() would unregister the parent loop's wakeup pipe from it.
+    global _loop, _thread, _lock
+    _loop = None
+    _thread = None
+    _lock = threading.Lock()
+
+
+if hasattr(os, "register_at_fork"):
+    os.register_at_fork(after_in_child=_reset_after_fork)
+
+
 def ensure_async_engine_loop() -> asyncio.AbstractEventLoop:
     """Get or create a persistent event loop for async engine work.
 

{data_designer_engine-0.9.6 → data_designer_engine-0.9.7}/src/data_designer/engine/mcp/io.py

@@ -32,6 +32,7 @@ import asyncio
 import atexit
 import json
 import logging
+import os
 import re
 import threading
 from collections.abc import Callable, Coroutine, Iterable
@@ -393,6 +394,14 @@ class MCPIOService:
     def _all_tools_keys(self) -> set[str]:
         return set(self._tools_cache) | set(self._inflight_tools) | set(self._tools_cache_epoch)
 
+    def _reset_after_fork(self) -> None:
+        # Same as async_concurrency._reset_after_fork: the loop thread did not survive
+        # the fork, and every session and in-flight task is bound to that dead loop.
+        self._loop = None
+        self._thread = None
+        self._loop_lock = threading.Lock()
+        self._reset_state()
+
     def _reset_state(self) -> None:
         self._sessions.clear()
         self._session_contexts.clear()
@@ -404,6 +413,8 @@ class MCPIOService:
 
 _MCP_IO_SERVICE = MCPIOService()
 atexit.register(_MCP_IO_SERVICE.shutdown)
+if hasattr(os, "register_at_fork"):
+    os.register_at_fork(after_in_child=_MCP_IO_SERVICE._reset_after_fork)
 
 
 def list_tools(provider: MCPProviderT, timeout_sec: float | None = None) -> tuple[MCPToolDefinition, ...]:

{data_designer_engine-0.9.6 → data_designer_engine-0.9.7}/src/data_designer/engine/sampling_gen/constraints.py

@@ -6,8 +6,6 @@ from __future__ import annotations
 from abc import ABC, abstractmethod
 from typing import TYPE_CHECKING
 
-from numpy.typing import NDArray
-
 import data_designer.lazy_heavy_imports as lazy
 from data_designer.config.base import ConfigBase
 from data_designer.config.sampler_constraints import (
@@ -21,6 +19,7 @@ from data_designer.config.sampler_constraints import (
 if TYPE_CHECKING:
     import numpy as np
     import pandas as pd
+    from numpy.typing import NDArray
 
 
 class ConstraintChecker(ConfigBase, ABC):

data_designer_engine-0.9.7/tests/engine/dataset_builders/utils/test_async_concurrency.py

@@ -0,0 +1,45 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+from __future__ import annotations
+
+from tests.engine.fork_probe import requires_fork, run_fork_probe
+
+
+@requires_fork
+def test_forked_child_starts_fresh_loop_after_parent_used_it() -> None:
+    run_fork_probe(
+        """
+        from data_designer.engine.dataset_builders.utils import async_concurrency
+
+
+        def child(queue):
+            loop = async_concurrency.ensure_async_engine_loop()
+            queue.put(asyncio.run_coroutine_threadsafe(ping(), loop).result(timeout=5))
+
+
+        parent_loop = async_concurrency.ensure_async_engine_loop()
+        assert asyncio.run_coroutine_threadsafe(ping(), parent_loop).result(timeout=5) == "pong"
+        assert run_forked_child(child) == "pong"
+        assert async_concurrency.ensure_async_engine_loop() is parent_loop
+        assert asyncio.run_coroutine_threadsafe(ping(), parent_loop).result(timeout=5) == "pong"
+        """
+    )
+
+
+@requires_fork
+def test_forked_child_does_not_inherit_held_loop_lock() -> None:
+    run_fork_probe(
+        """
+        from data_designer.engine.dataset_builders.utils import async_concurrency
+
+
+        def child(queue):
+            loop = async_concurrency.ensure_async_engine_loop()
+            queue.put(asyncio.run_coroutine_threadsafe(ping(), loop).result(timeout=5))
+
+
+        with async_concurrency._lock:
+            assert run_forked_child(child) == "pong"
+        """
+    )

data_designer_engine-0.9.7/tests/engine/fork_probe.py

@@ -0,0 +1,62 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+"""Run fork-safety scenarios in a fresh interpreter.
+
+Forking the pytest worker copies everything earlier tests left behind (DuckDB
+thread pools, large allocator reservations), which makes fork() itself fail
+with ENOMEM on memory-constrained CI runners. A fresh interpreter gives each
+scenario a small, clean parent.
+"""
+
+from __future__ import annotations
+
+import multiprocessing
+import subprocess
+import sys
+import textwrap
+
+import pytest
+
+requires_fork = pytest.mark.skipif(
+    "fork" not in multiprocessing.get_all_start_methods(), reason="requires the fork start method"
+)
+
+_CHILD_JOIN_TIMEOUT_S = 20
+_SCRIPT_TIMEOUT_S = 90
+
+_PRELUDE = f"""
+import asyncio
+import multiprocessing
+import sys
+
+
+async def ping():
+    return "pong"
+
+
+def run_forked_child(target):
+    ctx = multiprocessing.get_context("fork")
+    queue = ctx.Queue()
+    child = ctx.Process(target=target, args=(queue,))
+    child.start()
+    child.join(timeout={_CHILD_JOIN_TIMEOUT_S})
+    if child.is_alive():
+        child.kill()
+        child.join()
+        sys.exit("forked child hung")
+    if child.exitcode != 0:
+        sys.exit(f"forked child exited with code {{child.exitcode}}")
+    return queue.get(timeout=5)
+"""
+
+
+def run_fork_probe(body: str) -> None:
+    """Run ``body`` after a prelude that defines ``ping()`` and ``run_forked_child(target)``.
+
+    ``target`` receives a queue and must put its result on it. The probe fails if
+    the script exits non-zero, which includes a forked child that hangs.
+    """
+    script = _PRELUDE + textwrap.dedent(body)
+    result = subprocess.run([sys.executable, "-c", script], capture_output=True, text=True, timeout=_SCRIPT_TIMEOUT_S)
+    assert result.returncode == 0, f"fork probe failed:\n{result.stderr}"

{data_designer_engine-0.9.6 → data_designer_engine-0.9.7}/tests/engine/mcp/test_mcp_io.py

@@ -12,6 +12,7 @@ from data_designer.config.mcp import LocalStdioMCPProvider, MCPProvider
 from data_designer.engine.mcp import io as mcp_io
 from data_designer.engine.mcp.errors import MCPToolError
 from data_designer.engine.mcp.registry import MCPToolDefinition, MCPToolResult
+from tests.engine.fork_probe import requires_fork, run_fork_probe
 
 
 @pytest.fixture(autouse=True)
@@ -1283,3 +1284,28 @@ def test_call_tools_timeout_error_without_timeout_value(
 
     with pytest.raises(MCPToolError, match="Timed out after unknown"):
         mcp_io.call_tools([(stub_stdio_provider, "test_tool", {})], timeout_sec=None)
+
+
+# =============================================================================
+# Fork safety tests
+# =============================================================================
+
+
+@requires_fork
+def test_forked_child_starts_fresh_mcp_loop_after_parent_used_it() -> None:
+    run_fork_probe(
+        """
+        from data_designer.engine.mcp import io as mcp_io
+
+        service = mcp_io._MCP_IO_SERVICE
+
+
+        def child(queue):
+            queue.put(service._run_on_loop(ping(), timeout_sec=5))
+
+
+        assert service._run_on_loop(ping(), timeout_sec=5) == "pong"
+        assert run_forked_child(child) == "pong"
+        assert service._run_on_loop(ping(), timeout_sec=5) == "pong"
+        """
+    )