feat(components): add a DYN_JEMALLOC preload opt-in for the router and indexer

components/src/dynamo/frontend/__main__.py

@@ -1,55 +1,11 @@
 #  SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 #  SPDX-License-Identifier: Apache-2.0
 
-import ctypes.util
-import os
-import sys
-
-
-def _maybe_preload_jemalloc() -> None:
-    """Enable jemalloc before importing the frontend runtime.
-
-    Set DYN_FRONTEND_JEMALLOC to 1, true, or yes (case-insensitive).
-    LD_PRELOAD takes effect at process start, so re-exec once. Child processes
-    inherit the preload. If the library is missing or exec fails, warn and
-    continue with the current allocator.
-    """
-    # Match env_bool without importing Dynamo utilities before the re-exec.
-    if os.environ.get("DYN_FRONTEND_JEMALLOC", "").lower() not in ("1", "true", "yes"):
-        return
-    existing = os.environ.get("LD_PRELOAD", "")
-    if any(
-        os.path.basename(entry).startswith("libjemalloc")
-        for entry in existing.replace(":", " ").split()
-    ):
-        return  # already configured (or we already re-exec'd)
-
-    lib = ctypes.util.find_library("jemalloc")
-    if not lib:
-        # Logging is not configured yet, so write the warning to stderr.
-        print(
-            "WARNING: DYN_FRONTEND_JEMALLOC is enabled but libjemalloc was not found "
-            "(install libjemalloc2); continuing with the default allocator.",
-            file=sys.stderr,
-        )
-        return
-
-    env = os.environ.copy()
-    env["LD_PRELOAD"] = f"{lib}:{existing}" if existing else lib
-    sys.stdout.flush()
-    sys.stderr.flush()
-    try:
-        os.execve(sys.executable, sys.orig_argv, env)
-    except OSError as exc:
-        print(
-            f"WARNING: Could not re-exec frontend with jemalloc: {exc}; "
-            "continuing with the current allocator.",
-            file=sys.stderr,
-        )
-
+from dynamo._jemalloc import maybe_preload_jemalloc
 
 if __name__ == "__main__":
-    _maybe_preload_jemalloc()
+    # DYN_FRONTEND_JEMALLOC predates DYN_JEMALLOC and remains a frontend-only alias.
+    maybe_preload_jemalloc(alias="DYN_FRONTEND_JEMALLOC")
 
     # Import the runtime only after configuring the allocator.
     from dynamo.frontend.main import main

components/src/dynamo/frontend/tests/test_jemalloc.py

@@ -1,6 +1,11 @@
 # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: Apache-2.0
 
+"""Frontend entry-point wiring for the jemalloc preload.
+
+lib/bindings/python/tests/test_jemalloc_preload.py covers the preload itself.
+"""
+
 import ctypes.util
 import os
 import runpy
@@ -18,116 +23,41 @@
 
 @pytest.fixture
 def startup(monkeypatch):
-    monkeypatch.delenv("DYN_FRONTEND_JEMALLOC", raising=False)
-    monkeypatch.delenv("LD_PRELOAD", raising=False)
+    for name in ("DYN_JEMALLOC", "DYN_FRONTEND_JEMALLOC", "LD_PRELOAD"):
+        monkeypatch.delenv(name, raising=False)
     runtime = ModuleType("dynamo.frontend.main")
     runtime.main = Mock()
     monkeypatch.setitem(sys.modules, "dynamo.frontend.main", runtime)
-    find_library = Mock(return_value="libjemalloc.so.2")
     execve = Mock(side_effect=SystemExit)
-    monkeypatch.setattr(ctypes.util, "find_library", find_library)
+    monkeypatch.setattr(
+        ctypes.util, "find_library", Mock(return_value="libjemalloc.so.2")
+    )
     monkeypatch.setattr(os, "execve", execve)
-    return find_library, execve, runtime.main
+    return execve, runtime.main
 
 
-@pytest.mark.parametrize("enabled", [None, "0", "false"])
-def test_disabled(startup, monkeypatch, enabled):
-    if enabled is not None:
-        monkeypatch.setenv("DYN_FRONTEND_JEMALLOC", enabled)
-    find_library, execve, main = startup
+def test_disabled(startup):
+    execve, main = startup
     runpy.run_path(str(ENTRYPOINT), run_name="__main__")
-    find_library.assert_not_called()
     execve.assert_not_called()
     main.assert_called_once_with()
-    assert "LD_PRELOAD" not in os.environ
 
 
-@pytest.mark.parametrize("separator", [":", " "])
-def test_already_preloaded(startup, monkeypatch, separator):
-    monkeypatch.setenv("DYN_FRONTEND_JEMALLOC", "1")
-    preload = f"libother.so{separator}/usr/lib/libjemalloc.so.2"
-    monkeypatch.setenv("LD_PRELOAD", preload)
-    find_library, execve, main = startup
-    runpy.run_path(str(ENTRYPOINT), run_name="__main__")
-    find_library.assert_not_called()
-    execve.assert_not_called()
-    main.assert_called_once_with()
-    assert os.environ["LD_PRELOAD"] == preload
-
-
-def test_missing_library(startup, monkeypatch, capsys):
-    monkeypatch.setenv("DYN_FRONTEND_JEMALLOC", "1")
-    find_library, execve, main = startup
-    find_library.return_value = None
-    runpy.run_path(str(ENTRYPOINT), run_name="__main__")
-    find_library.assert_called_once_with("jemalloc")
-    execve.assert_not_called()
-    main.assert_called_once_with()
-    assert "libjemalloc was not found" in capsys.readouterr().err
-    assert "LD_PRELOAD" not in os.environ
-
-
-@pytest.mark.parametrize(
-    "enabled,existing",
-    [
-        ("1", ""),
-        ("true", "libother.so:libanother.so"),
-        ("YES", "/opt/jemalloc-tools/libheaptrace.so"),
-    ],
-)
-def test_restart(startup, monkeypatch, enabled, existing):
-    monkeypatch.setenv("DYN_FRONTEND_JEMALLOC", enabled)
-    monkeypatch.setenv("LD_PRELOAD", existing)
-    argv = [
-        sys.executable,
-        "-u",
-        "-X",
-        "faulthandler",
-        "-m",
-        "dynamo.frontend",
-        "--http-port",
-        "8123",
-    ]
-    monkeypatch.setattr(sys, "orig_argv", argv)
-    find_library, execve, main = startup
-    stdout, stderr = Mock(), Mock()
-    monkeypatch.setattr(sys, "stdout", stdout)
-    monkeypatch.setattr(sys, "stderr", stderr)
-
-    def replace_process(*args):
-        stdout.flush.assert_called_once_with()
-        stderr.flush.assert_called_once_with()
-        raise SystemExit
-
-    execve.side_effect = replace_process
-    expected_env = dict(
-        os.environ, LD_PRELOAD="libjemalloc.so.2" + (f":{existing}" if existing else "")
-    )
+@pytest.mark.parametrize("env_var", ["DYN_JEMALLOC", "DYN_FRONTEND_JEMALLOC"])
+def test_preloads_before_runtime_import(startup, monkeypatch, env_var):
+    monkeypatch.setenv(env_var, "1")
+    # Importing the runtime before the re-exec raises ImportError, not SystemExit.
+    monkeypatch.setitem(sys.modules, "dynamo.frontend.main", None)
+    execve, _ = startup
     with pytest.raises(SystemExit):
         runpy.run_path(str(ENTRYPOINT), run_name="__main__")
-    find_library.assert_called_once_with("jemalloc")
-    execve.assert_called_once_with(sys.executable, argv, expected_env)
-    assert os.environ["LD_PRELOAD"] == existing
-    main.assert_not_called()
-
-
-@pytest.mark.parametrize("existing", [None, "", "libother.so"])
-def test_exec_failure(startup, monkeypatch, capsys, existing):
-    monkeypatch.setenv("DYN_FRONTEND_JEMALLOC", "1")
-    if existing is not None:
-        monkeypatch.setenv("LD_PRELOAD", existing)
-    _, execve, main = startup
-    execve.side_effect = OSError("exec blocked")
-    runpy.run_path(str(ENTRYPOINT), run_name="__main__")
-    main.assert_called_once_with()
-    assert os.environ.get("LD_PRELOAD") == existing
-    assert "exec blocked" in capsys.readouterr().err
+    _, _, env = execve.call_args.args
+    assert env["LD_PRELOAD"] == "libjemalloc.so.2"
 
 
 def test_import_does_not_restart(startup, monkeypatch):
-    monkeypatch.setenv("DYN_FRONTEND_JEMALLOC", "1")
-    find_library, execve, main = startup
+    monkeypatch.setenv("DYN_JEMALLOC", "1")
+    execve, main = startup
     runpy.run_path(str(ENTRYPOINT), run_name="dynamo.frontend.__main__")
-    find_library.assert_not_called()
     execve.assert_not_called()
     main.assert_not_called()

components/src/dynamo/router/AGENTS.md

@@ -98,7 +98,8 @@ logic but do not share the same serialization, RPC, or process boundaries.
 |---|---|
 | `../frontend/main.py` | Selects router mode and optional Python chat processor |
 | `../frontend/vllm_processor.py` | vLLM-native Python pre/postprocessor |
-| `__main__.py` | Standalone Python router service |
+| `__main__.py` | Entry point; applies the opt-in jemalloc preload before importing `main.py` |
+| `main.py` | Standalone Python router service |
 | `../thunderagent_router/__main__.py` | Custom registered router facade |
 | `../../../../lib/llm/src/entrypoint/input/common.rs` | Builds embedded Rust routing pipelines |
 | `../../../../lib/llm/src/kv_router/routing_host.rs` | Routed-generation orchestration and router-side tracker updates |