binaryaaron · Oct 22, 2026 · Oct 20, 2026 · Oct 21, 2026 · Oct 21, 2026 · Oct 21, 2026
@@ -26,7 +26,7 @@ Preparation (`_prepare_async_run`):
 4. Constructs `CompletionTracker`, `RowGroupBufferManager`, `AsyncTaskScheduler`
 5. Hooks `ProcessorRunner` for pre-batch and post-batch stages
 
-`AsyncTaskScheduler` runs on a dedicated async loop with frontier-driven dispatch, task-admission leases, salvage rounds for failed tasks, and order-dependent locks for columns that must execute sequentially. Ready frontier tasks enter `FairTaskQueue`, are selected through virtual-time ordering, and are committed only after `TaskAdmissionController` acquires the required scheduler resources. Salvage-exhausted tasks are dropped except for preserved retryable failures: provider rate limits and local request-admission queue timeouts stay deferred and retry after cooldown/backoff so scheduler-local pressure delays records rather than discarding them.
+`AsyncTaskScheduler` runs on a dedicated async loop (a process-wide singleton from `ensure_async_engine_loop()`; a forked child drops the inherited loop and starts its own) with frontier-driven dispatch, task-admission leases, salvage rounds for failed tasks, and order-dependent locks for columns that must execute sequentially. Ready frontier tasks enter `FairTaskQueue`, are selected through virtual-time ordering, and are committed only after `TaskAdmissionController` acquires the required scheduler resources. Salvage-exhausted tasks are dropped except for preserved retryable failures: provider rate limits and local request-admission queue timeouts stay deferred and retry after cooldown/backoff so scheduler-local pressure delays records rather than discarding them.
 
 ### Execution Graph
 

@@ -55,7 +55,7 @@ Maps `tool_alias` → `ToolConfig`. Lazy `MCPFacade` construction mirrors `Model
 
 ## Design Decisions
 
-- **Single background async loop** avoids creating event loops per request. All MCP I/O funnels through one loop on a daemon thread, with sync callers bridging via `run_coroutine_threadsafe`.
+- **Single background async loop** avoids creating event loops per request. All MCP I/O funnels through one loop on a daemon thread, with sync callers bridging via `run_coroutine_threadsafe`. An `os.register_at_fork` hook drops the loop, thread, lock and sessions in a forked child, because the loop's thread does not survive `fork()`; the child starts a fresh loop on first use. The hook never closes the inherited loop, which would unregister the parent loop's wakeup pipe from the epoll instance they share.
 - **Session pooling with in-flight deduplication** prevents redundant connections when multiple generators discover tools from the same provider concurrently.
 - **Tool schema coalescing** — concurrent `list_tools` calls for the same session share a single in-flight request, reducing startup latency when many columns use the same tool.
 - **Turn limits on tool loops** prevent runaway tool-call chains. Past the limit, `refuse_completion_response` answers tool calls with a refusal, and the loop ends once the model replies without calling a tool. That relies on the model being free not to call one, which is why a caller's forced `tool_choice` is dropped after the first tool call. A `tool_choice` inside `extra_body` is sent unchanged on every round.

@@ -19,12 +19,19 @@
     ``_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.
 

@@ -32,6 +32,7 @@
 import atexit
 import json
 import logging
+import os
 import re
 import threading
 from collections.abc import Callable, Coroutine, Iterable
@@ -393,6 +394,14 @@ def _invalidate_tools_cache(self, keys: Iterable[str]) -> None:
     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 @@ def _reset_state(self) -> None:
 
 _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, ...]:

@@ -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"
+        """
+    )
@@ -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}"
@@ -12,6 +12,7 @@
 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 mock_run_coroutine_threadsafe(coro: Any, loop: Any) -> MockFuture:
 
     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"
+        """
+    )
@@ -0,0 +1,56 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+from __future__ import annotations
+
+import multiprocessing
+import os
+import subprocess
+import sys
+
+import pytest
+
+# Runs in a fresh interpreter so the forked parent is small and holds only what the
+# scenario creates, not whatever earlier tests left in the pytest worker.
+_PREVIEW_IN_FORKED_CHILD = """
+import multiprocessing
+import sys
+
+import data_designer.config as dd
+from data_designer.interface import DataDesigner
+
+
+def preview():
+    builder = dd.DataDesignerConfigBuilder()
+    builder.add_column(
+        dd.SamplerColumnConfig(
+            name="category",
+            sampler_type=dd.SamplerType.CATEGORY,
+            params=dd.CategorySamplerParams(values=["a", "b", "c"]),
+        )
+    )
+    assert len(DataDesigner().preview(builder, num_records=3).dataset) == 3
+
+
+preview()  # the parent starts DataDesigner's background event loop
+child = multiprocessing.get_context("fork").Process(target=preview)
+child.start()
+child.join(timeout=30)
+if child.is_alive():
+    child.kill()
+    child.join()
+    sys.exit("forked child hung in preview()")
+sys.exit(child.exitcode)
+"""
+
+
+@pytest.mark.skipif("fork" not in multiprocessing.get_all_start_methods(), reason="requires the fork start method")
+def test_preview_in_child_forked_after_parent_preview() -> None:
+    result = subprocess.run(
+        [sys.executable, "-c", _PREVIEW_IN_FORKED_CHILD],
+        capture_output=True,
+        text=True,
+        timeout=120,
+        env={**os.environ, "NEMO_TELEMETRY_ENABLED": "false"},
+    )
+    assert result.returncode == 0, f"preview() in a forked child failed:\n{result.stderr[-4000:]}"