diff --git a/architecture/dataset-builders.md b/architecture/dataset-builders.md index 88d4f6a0b..3738f31ce 100644 --- a/architecture/dataset-builders.md +++ b/architecture/dataset-builders.md @@ -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 diff --git a/architecture/mcp.md b/architecture/mcp.md index 4024cabc5..c7daefd9f 100644 --- a/architecture/mcp.md +++ b/architecture/mcp.md @@ -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. diff --git a/packages/data-designer-engine/src/data_designer/engine/dataset_builders/utils/async_concurrency.py b/packages/data-designer-engine/src/data_designer/engine/dataset_builders/utils/async_concurrency.py index eb9e2104a..751333110 100644 --- a/packages/data-designer-engine/src/data_designer/engine/dataset_builders/utils/async_concurrency.py +++ b/packages/data-designer-engine/src/data_designer/engine/dataset_builders/utils/async_concurrency.py @@ -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. diff --git a/packages/data-designer-engine/src/data_designer/engine/mcp/io.py b/packages/data-designer-engine/src/data_designer/engine/mcp/io.py index d9b2b8efd..bc188fe36 100644 --- a/packages/data-designer-engine/src/data_designer/engine/mcp/io.py +++ b/packages/data-designer-engine/src/data_designer/engine/mcp/io.py @@ -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, ...]: diff --git a/packages/data-designer-engine/tests/engine/dataset_builders/utils/test_async_concurrency.py b/packages/data-designer-engine/tests/engine/dataset_builders/utils/test_async_concurrency.py new file mode 100644 index 000000000..80627ca41 --- /dev/null +++ b/packages/data-designer-engine/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" + """ + ) diff --git a/packages/data-designer-engine/tests/engine/fork_probe.py b/packages/data-designer-engine/tests/engine/fork_probe.py new file mode 100644 index 000000000..14501da33 --- /dev/null +++ b/packages/data-designer-engine/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}" diff --git a/packages/data-designer-engine/tests/engine/mcp/test_mcp_io.py b/packages/data-designer-engine/tests/engine/mcp/test_mcp_io.py index adea98441..65eb2c246 100644 --- a/packages/data-designer-engine/tests/engine/mcp/test_mcp_io.py +++ b/packages/data-designer-engine/tests/engine/mcp/test_mcp_io.py @@ -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" + """ + ) diff --git a/packages/data-designer/tests/interface/test_fork_safety.py b/packages/data-designer/tests/interface/test_fork_safety.py new file mode 100644 index 000000000..87dde0102 --- /dev/null +++ b/packages/data-designer/tests/interface/test_fork_safety.py @@ -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:]}"