Skip to content

Commit 4ed3a9e

Browse files
authored
fix: reset background event loops in forked child processes (#1069)
Drop the inherited engine and MCP event loops in a forked child so its first run starts fresh loops instead of hanging. Fixes #1068. Signed-off-by: Aaron Gonzales <aagonzales@nvidia.com>
1 parent 9a09fce commit 4ed3a9e

8 files changed

Lines changed: 222 additions & 2 deletions

File tree

‎architecture/dataset-builders.md‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -26,7 +26,7 @@ Preparation (`_prepare_async_run`):
2626
4. Constructs `CompletionTracker`, `RowGroupBufferManager`, `AsyncTaskScheduler`
2727
5. Hooks `ProcessorRunner` for pre-batch and post-batch stages
2828

29-
`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.
29+
`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.
3030

3131
### Execution Graph
3232

‎architecture/mcp.md‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -55,7 +55,7 @@ Maps `tool_alias` → `ToolConfig`. Lazy `MCPFacade` construction mirrors `Model
5555

5656
## Design Decisions
5757

58-
- **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`.
58+
- **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.
5959
- **Session pooling with in-flight deduplication** prevents redundant connections when multiple generators discover tools from the same provider concurrently.
6060
- **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.
6161
- **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.

‎packages/data-designer-engine/src/data_designer/engine/dataset_builders/utils/async_concurrency.py‎

Lines changed: 20 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -19,12 +19,19 @@
1919
``_loop.is_running() == False`` before the first loop has entered
2020
``run_forever()``, which would create a duplicate loop. On timeout,
2121
globals are reset and the orphaned loop is cleaned up before raising.
22+
23+
Fork safety:
24+
Only the forking thread survives ``os.fork()``, but the copied loop
25+
object still reports ``is_running() == True``. An ``after_in_child``
26+
hook drops the inherited loop, thread, and lock so the child's first
27+
call starts a fresh loop instead of waiting on one nothing drives.
2228
"""
2329

2430
from __future__ import annotations
2531

2632
import asyncio
2733
import logging
34+
import os
2835
import threading
2936

3037
logger = logging.getLogger(__name__)
@@ -42,6 +49,19 @@ def _run_loop(loop: asyncio.AbstractEventLoop, ready: threading.Event) -> None:
4249
loop.run_forever()
4350

4451

52+
def _reset_after_fork() -> None:
53+
# Drop, never close(), the inherited loop: the child shares the parent's epoll
54+
# instance, and close() would unregister the parent loop's wakeup pipe from it.
55+
global _loop, _thread, _lock
56+
_loop = None
57+
_thread = None
58+
_lock = threading.Lock()
59+
60+
61+
if hasattr(os, "register_at_fork"):
62+
os.register_at_fork(after_in_child=_reset_after_fork)
63+
64+
4565
def ensure_async_engine_loop() -> asyncio.AbstractEventLoop:
4666
"""Get or create a persistent event loop for async engine work.
4767

‎packages/data-designer-engine/src/data_designer/engine/mcp/io.py‎

Lines changed: 11 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -32,6 +32,7 @@
3232
import atexit
3333
import json
3434
import logging
35+
import os
3536
import re
3637
import threading
3738
from collections.abc import Callable, Coroutine, Iterable
@@ -393,6 +394,14 @@ def _invalidate_tools_cache(self, keys: Iterable[str]) -> None:
393394
def _all_tools_keys(self) -> set[str]:
394395
return set(self._tools_cache) | set(self._inflight_tools) | set(self._tools_cache_epoch)
395396

397+
def _reset_after_fork(self) -> None:
398+
# Same as async_concurrency._reset_after_fork: the loop thread did not survive
399+
# the fork, and every session and in-flight task is bound to that dead loop.
400+
self._loop = None
401+
self._thread = None
402+
self._loop_lock = threading.Lock()
403+
self._reset_state()
404+
396405
def _reset_state(self) -> None:
397406
self._sessions.clear()
398407
self._session_contexts.clear()
@@ -404,6 +413,8 @@ def _reset_state(self) -> None:
404413

405414
_MCP_IO_SERVICE = MCPIOService()
406415
atexit.register(_MCP_IO_SERVICE.shutdown)
416+
if hasattr(os, "register_at_fork"):
417+
os.register_at_fork(after_in_child=_MCP_IO_SERVICE._reset_after_fork)
407418

408419

409420
def list_tools(provider: MCPProviderT, timeout_sec: float | None = None) -> tuple[MCPToolDefinition, ...]:
Lines changed: 45 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,45 @@
1+
# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
2+
# SPDX-License-Identifier: Apache-2.0
3+
4+
from __future__ import annotations
5+
6+
from tests.engine.fork_probe import requires_fork, run_fork_probe
7+
8+
9+
@requires_fork
10+
def test_forked_child_starts_fresh_loop_after_parent_used_it() -> None:
11+
run_fork_probe(
12+
"""
13+
from data_designer.engine.dataset_builders.utils import async_concurrency
14+
15+
16+
def child(queue):
17+
loop = async_concurrency.ensure_async_engine_loop()
18+
queue.put(asyncio.run_coroutine_threadsafe(ping(), loop).result(timeout=5))
19+
20+
21+
parent_loop = async_concurrency.ensure_async_engine_loop()
22+
assert asyncio.run_coroutine_threadsafe(ping(), parent_loop).result(timeout=5) == "pong"
23+
assert run_forked_child(child) == "pong"
24+
assert async_concurrency.ensure_async_engine_loop() is parent_loop
25+
assert asyncio.run_coroutine_threadsafe(ping(), parent_loop).result(timeout=5) == "pong"
26+
"""
27+
)
28+
29+
30+
@requires_fork
31+
def test_forked_child_does_not_inherit_held_loop_lock() -> None:
32+
run_fork_probe(
33+
"""
34+
from data_designer.engine.dataset_builders.utils import async_concurrency
35+
36+
37+
def child(queue):
38+
loop = async_concurrency.ensure_async_engine_loop()
39+
queue.put(asyncio.run_coroutine_threadsafe(ping(), loop).result(timeout=5))
40+
41+
42+
with async_concurrency._lock:
43+
assert run_forked_child(child) == "pong"
44+
"""
45+
)
Lines changed: 62 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,62 @@
1+
# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
2+
# SPDX-License-Identifier: Apache-2.0
3+
4+
"""Run fork-safety scenarios in a fresh interpreter.
5+
6+
Forking the pytest worker copies everything earlier tests left behind (DuckDB
7+
thread pools, large allocator reservations), which makes fork() itself fail
8+
with ENOMEM on memory-constrained CI runners. A fresh interpreter gives each
9+
scenario a small, clean parent.
10+
"""
11+
12+
from __future__ import annotations
13+
14+
import multiprocessing
15+
import subprocess
16+
import sys
17+
import textwrap
18+
19+
import pytest
20+
21+
requires_fork = pytest.mark.skipif(
22+
"fork" not in multiprocessing.get_all_start_methods(), reason="requires the fork start method"
23+
)
24+
25+
_CHILD_JOIN_TIMEOUT_S = 20
26+
_SCRIPT_TIMEOUT_S = 90
27+
28+
_PRELUDE = f"""
29+
import asyncio
30+
import multiprocessing
31+
import sys
32+
33+
34+
async def ping():
35+
return "pong"
36+
37+
38+
def run_forked_child(target):
39+
ctx = multiprocessing.get_context("fork")
40+
queue = ctx.Queue()
41+
child = ctx.Process(target=target, args=(queue,))
42+
child.start()
43+
child.join(timeout={_CHILD_JOIN_TIMEOUT_S})
44+
if child.is_alive():
45+
child.kill()
46+
child.join()
47+
sys.exit("forked child hung")
48+
if child.exitcode != 0:
49+
sys.exit(f"forked child exited with code {{child.exitcode}}")
50+
return queue.get(timeout=5)
51+
"""
52+
53+
54+
def run_fork_probe(body: str) -> None:
55+
"""Run ``body`` after a prelude that defines ``ping()`` and ``run_forked_child(target)``.
56+
57+
``target`` receives a queue and must put its result on it. The probe fails if
58+
the script exits non-zero, which includes a forked child that hangs.
59+
"""
60+
script = _PRELUDE + textwrap.dedent(body)
61+
result = subprocess.run([sys.executable, "-c", script], capture_output=True, text=True, timeout=_SCRIPT_TIMEOUT_S)
62+
assert result.returncode == 0, f"fork probe failed:\n{result.stderr}"

‎packages/data-designer-engine/tests/engine/mcp/test_mcp_io.py‎

Lines changed: 26 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -12,6 +12,7 @@
1212
from data_designer.engine.mcp import io as mcp_io
1313
from data_designer.engine.mcp.errors import MCPToolError
1414
from data_designer.engine.mcp.registry import MCPToolDefinition, MCPToolResult
15+
from tests.engine.fork_probe import requires_fork, run_fork_probe
1516

1617

1718
@pytest.fixture(autouse=True)
@@ -1283,3 +1284,28 @@ def mock_run_coroutine_threadsafe(coro: Any, loop: Any) -> MockFuture:
12831284

12841285
with pytest.raises(MCPToolError, match="Timed out after unknown"):
12851286
mcp_io.call_tools([(stub_stdio_provider, "test_tool", {})], timeout_sec=None)
1287+
1288+
1289+
# =============================================================================
1290+
# Fork safety tests
1291+
# =============================================================================
1292+
1293+
1294+
@requires_fork
1295+
def test_forked_child_starts_fresh_mcp_loop_after_parent_used_it() -> None:
1296+
run_fork_probe(
1297+
"""
1298+
from data_designer.engine.mcp import io as mcp_io
1299+
1300+
service = mcp_io._MCP_IO_SERVICE
1301+
1302+
1303+
def child(queue):
1304+
queue.put(service._run_on_loop(ping(), timeout_sec=5))
1305+
1306+
1307+
assert service._run_on_loop(ping(), timeout_sec=5) == "pong"
1308+
assert run_forked_child(child) == "pong"
1309+
assert service._run_on_loop(ping(), timeout_sec=5) == "pong"
1310+
"""
1311+
)
Lines changed: 56 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,56 @@
1+
# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
2+
# SPDX-License-Identifier: Apache-2.0
3+
4+
from __future__ import annotations
5+
6+
import multiprocessing
7+
import os
8+
import subprocess
9+
import sys
10+
11+
import pytest
12+
13+
# Runs in a fresh interpreter so the forked parent is small and holds only what the
14+
# scenario creates, not whatever earlier tests left in the pytest worker.
15+
_PREVIEW_IN_FORKED_CHILD = """
16+
import multiprocessing
17+
import sys
18+
19+
import data_designer.config as dd
20+
from data_designer.interface import DataDesigner
21+
22+
23+
def preview():
24+
builder = dd.DataDesignerConfigBuilder()
25+
builder.add_column(
26+
dd.SamplerColumnConfig(
27+
name="category",
28+
sampler_type=dd.SamplerType.CATEGORY,
29+
params=dd.CategorySamplerParams(values=["a", "b", "c"]),
30+
)
31+
)
32+
assert len(DataDesigner().preview(builder, num_records=3).dataset) == 3
33+
34+
35+
preview() # the parent starts DataDesigner's background event loop
36+
child = multiprocessing.get_context("fork").Process(target=preview)
37+
child.start()
38+
child.join(timeout=30)
39+
if child.is_alive():
40+
child.kill()
41+
child.join()
42+
sys.exit("forked child hung in preview()")
43+
sys.exit(child.exitcode)
44+
"""
45+
46+
47+
@pytest.mark.skipif("fork" not in multiprocessing.get_all_start_methods(), reason="requires the fork start method")
48+
def test_preview_in_child_forked_after_parent_preview() -> None:
49+
result = subprocess.run(
50+
[sys.executable, "-c", _PREVIEW_IN_FORKED_CHILD],
51+
capture_output=True,
52+
text=True,
53+
timeout=120,
54+
env={**os.environ, "NEMO_TELEMETRY_ENABLED": "false"},
55+
)
56+
assert result.returncode == 0, f"preview() in a forked child failed:\n{result.stderr[-4000:]}"

0 commit comments

Comments
 (0)