From f8521df4e561bb367f5ae327ffcb3ad435bb0147 Mon Sep 17 00:00:00 2001 From: PeaBrane Date: Wed, 7 Oct 2026 21:42:31 -0700 Subject: [PATCH 1/2] feat(components): add a DYN_JEMALLOC preload opt-in for the router and indexer Only `python -m dynamo.frontend` can preload jemalloc today, through DYN_FRONTEND_JEMALLOC. The standalone router and the standalone KV indexer have no opt-in and need a hand-written LD_PRELOAD with the library path. Move the frontend's preload helper into the stdlib-only module dynamo._jemalloc in the runtime wheel, so every entry point can call it before importing the Rust extension, including dynamo.indexer, which ships in that wheel. The helper reads the new DYN_JEMALLOC. An entry point can also pass a component-specific alias; only the frontend passes DYN_FRONTEND_JEMALLOC, and it preloads when either variable is enabled. Truthy parsing, already-preloaded detection, the missing-library warning, and the single re-exec are unchanged. Warnings now name the variable that enabled the preload. Call the helper before the runtime import in dynamo.frontend, dynamo.router, and dynamo.indexer. The router's module body moves verbatim from __main__.py to main.py, as in the frontend, so `python -m dynamo.router` keeps the same CLI. jemalloc stays opt-in. Signed-off-by: PeaBrane --- components/src/dynamo/frontend/__main__.py | 50 +--- .../dynamo/frontend/tests/test_jemalloc.py | 116 ++------- components/src/dynamo/router/AGENTS.md | 3 +- components/src/dynamo/router/__main__.py | 243 +---------------- components/src/dynamo/router/main.py | 245 ++++++++++++++++++ .../router/tests/test_router_jemalloc.py | 55 ++++ .../router/tests/test_standalone_router.py | 2 +- .../router/standalone-indexer.md | 8 + .../components/frontend-configuration.mdx | 16 ++ lib/bindings/python/src/dynamo/_jemalloc.py | 64 +++++ .../python/src/dynamo/indexer/__main__.py | 7 +- .../python/tests/test_jemalloc_preload.py | 162 ++++++++++++ 12 files changed, 590 insertions(+), 381 deletions(-) create mode 100644 components/src/dynamo/router/main.py create mode 100644 components/src/dynamo/router/tests/test_router_jemalloc.py create mode 100644 lib/bindings/python/src/dynamo/_jemalloc.py create mode 100644 lib/bindings/python/tests/test_jemalloc_preload.py diff --git a/components/src/dynamo/frontend/__main__.py b/components/src/dynamo/frontend/__main__.py index dc9bf39eb4f0..3095ad43ccdb 100644 --- a/components/src/dynamo/frontend/__main__.py +++ b/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 diff --git a/components/src/dynamo/frontend/tests/test_jemalloc.py b/components/src/dynamo/frontend/tests/test_jemalloc.py index 42676011ce67..93e61a4ad094 100644 --- a/components/src/dynamo/frontend/tests/test_jemalloc.py +++ b/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() diff --git a/components/src/dynamo/router/AGENTS.md b/components/src/dynamo/router/AGENTS.md index fb5b99f73b5b..ce25444d45c1 100644 --- a/components/src/dynamo/router/AGENTS.md +++ b/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 | diff --git a/components/src/dynamo/router/__main__.py b/components/src/dynamo/router/__main__.py index 53a9e4455c4d..d0fe9a0210db 100644 --- a/components/src/dynamo/router/__main__.py +++ b/components/src/dynamo/router/__main__.py @@ -1,245 +1,12 @@ # SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 -""" -Standalone KV Router Service +from dynamo._jemalloc import maybe_preload_jemalloc -Usage: python -m dynamo.router --endpoint [args] - -This service provides a standalone KV-aware router for any set of workers -in a Dynamo deployment. It can be used for disaggregated serving (e.g., routing -to prefill workers) or any other scenario requiring intelligent KV cache-aware -routing decisions. -""" - -import asyncio -import logging -from typing import Optional - -import uvloop - -from dynamo.common.utils.token_ids import token_ids_to_list -from dynamo.llm import AisPerfConfig, KvRouter, KvRouterConfig -from dynamo.router.args import ( - DynamoRouterConfig, - build_ais_perf_config, - build_kv_router_config, -) -from dynamo.router.args import parse_args as parse_router_args -from dynamo.runtime import Client, DistributedRuntime, dynamo_worker -from dynamo.runtime.logging import configure_dynamo_logging - -configure_dynamo_logging() -logger = logging.getLogger(__name__) - - -class StandaloneRouterHandler: - """Handles routing requests to workers using KV-aware routing.""" - - def __init__( - self, - runtime: DistributedRuntime, - worker_endpoint_path: str, - block_size: int, - kv_router_config: KvRouterConfig, - ais_perf_config: Optional[AisPerfConfig], - ): - self.runtime = runtime - self.worker_endpoint_path = worker_endpoint_path - self.block_size = block_size - self.kv_router_config = kv_router_config - self.ais_perf_config = ais_perf_config - self.kv_router: Optional[KvRouter] = None - self.worker_client: Optional[Client] = None - - async def initialize(self): - """Initialize the KV router for workers.""" - try: - # Parse endpoint path (format: namespace.component.endpoint) - parts = self.worker_endpoint_path.split(".") - if len(parts) != 3: - raise ValueError( - f"Invalid endpoint path format: {self.worker_endpoint_path}. " - "Expected format: namespace.component.endpoint" - ) - namespace, component, endpoint = parts - - # Get worker endpoint - worker_endpoint = self.runtime.endpoint( - f"{namespace}.{component}.{endpoint}" - ) - self.worker_client = await worker_endpoint.client() - - self.kv_router = KvRouter( - endpoint=worker_endpoint, - block_size=self.block_size, - kv_router_config=self.kv_router_config, - ais_perf_config=self.ais_perf_config, - ) - - except Exception as e: - logger.error(f"Failed to initialize KvRouter: {e}") - raise - - async def generate(self, request): - """ - Generate tokens using the KV-aware router. - - Routes a PreprocessedRequest-shaped request to the best worker and streams back - its LLMEngineOutput responses unchanged, so fields this router does not inspect - (multimodal data, KV hints, agent context, ...) survive the hop. - """ - if self.kv_router is None: - logger.error("KvRouter not initialized - cannot process request") - raise RuntimeError("Router not initialized") - - preprocessed_request = dict(request) - preprocessed_request.setdefault("model", "unknown") - # Legacy callers send a top-level dp_rank instead of routing hints. - dp_rank = preprocessed_request.get("dp_rank") - if preprocessed_request.get("routing") is None and dp_rank is not None: - preprocessed_request["routing"] = {"dp_rank": dp_rank} - - async for worker_output in await self.kv_router.generate_from_request( - preprocessed_request - ): - yield worker_output - - async def best_worker_id( - self, token_ids, router_config_override=None, cache_namespace=None - ): - """ - Get the best worker ID for a given set of tokens without actually routing. - - This method returns the worker ID that would be selected based on KV cache - overlap, but does NOT actually route the request or update router states. - It's useful for debugging, monitoring, or implementing custom routing logic. - """ - if self.kv_router is None: - logger.error("KvRouter not initialized - cannot get best worker") - raise RuntimeError("Router not initialized") - - (worker_id, _dp_rank, _overlap_blocks) = await self.kv_router.best_worker( - token_ids, - router_config_override, - cache_namespace=cache_namespace, - ) - - yield worker_id - - async def get_overlap_scores(self, request): - """ - Get per-worker KV overlap by storage tier without routing the request. - - This endpoint returns matched blocks for each worker_id/dp_rank pair. - Shared-cache hits are request-global and are also reported per row as - blocks beyond that rank's device-local prefix. - """ - if self.kv_router is None: - logger.error("KvRouter not initialized - cannot get overlap scores") - raise RuntimeError("Router not initialized") - - scores = await self.kv_router.get_overlap_scores( - token_ids_to_list(request["token_ids"]), - request.get("router_config_override"), - request.get("block_mm_infos"), - request.get("lora_name"), - request.get("include_shared", True), - request.get("cache_namespace"), - ) - - yield scores - - -def parse_args(argv=None) -> DynamoRouterConfig: - """Parse router CLI arguments (compatibility shim delegating to args.parse_args).""" - return parse_router_args(argv) - - -@dynamo_worker() -async def worker(runtime: DistributedRuntime): - """Main worker function for the standalone router service.""" - - config = parse_args() - - logger.info("Starting Standalone Router Service") - logger.debug( - "Configuration: endpoint=%s, router_block_size=%s, " - "overlap_score_credit=%s, overlap_score_credit_decay=%s, " - "prefill_load_scale=%s, decode_active_request_weight=%s, " - "router_temperature=%s, use_kv_events=%s, router_replica_sync=%s, " - "router_track_active_blocks=%s, router_track_output_blocks=%s, " - "router_assume_kv_reuse=%s, router_track_prefill_tokens=%s, " - "router_ttl_secs=%s, router_approximate_cache_policy=%s", - config.endpoint, - config.router_block_size, - config.overlap_score_credit, - config.overlap_score_credit_decay, - config.prefill_load_scale, - config.decode_active_request_weight, - config.router_temperature, - config.use_kv_events, - config.router_replica_sync, - config.router_track_active_blocks, - config.router_track_output_blocks, - config.router_assume_kv_reuse, - config.router_track_prefill_tokens, - config.router_ttl_secs, - config.router_approximate_cache_policy, - ) - - kv_router_config = build_kv_router_config(config) - ais_perf_config = build_ais_perf_config(config) - - # Create handler - handler = StandaloneRouterHandler( - runtime, - config.endpoint, - config.router_block_size, - kv_router_config, - ais_perf_config, - ) - await handler.initialize() - - # Create endpoints - generate_endpoint = runtime.endpoint(f"{config.namespace}.router.generate") - best_worker_endpoint = runtime.endpoint(f"{config.namespace}.router.best_worker_id") - overlap_scores_endpoint = runtime.endpoint( - f"{config.namespace}.router.get_overlap_scores" - ) - - logger.debug("Starting to serve endpoints...") - - # Serve both endpoints concurrently - try: - await asyncio.gather( - generate_endpoint.serve_endpoint( - handler.generate, - graceful_shutdown=True, - metrics_labels=[("service", "router")], - ), - best_worker_endpoint.serve_endpoint( - handler.best_worker_id, - graceful_shutdown=True, - metrics_labels=[("service", "router")], - ), - overlap_scores_endpoint.serve_endpoint( - handler.get_overlap_scores, - graceful_shutdown=True, - metrics_labels=[("service", "router")], - ), - ) - except Exception as e: - logger.error(f"Failed to serve endpoint: {e}") - raise - finally: - logger.info("Standalone Router Service shutting down") - - -def main(): - """Entry point for the standalone router service.""" - uvloop.run(worker()) +if __name__ == "__main__": + maybe_preload_jemalloc() + # Import the runtime only after configuring the allocator. + from dynamo.router.main import main -if __name__ == "__main__": main() diff --git a/components/src/dynamo/router/main.py b/components/src/dynamo/router/main.py new file mode 100644 index 000000000000..53a9e4455c4d --- /dev/null +++ b/components/src/dynamo/router/main.py @@ -0,0 +1,245 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +""" +Standalone KV Router Service + +Usage: python -m dynamo.router --endpoint [args] + +This service provides a standalone KV-aware router for any set of workers +in a Dynamo deployment. It can be used for disaggregated serving (e.g., routing +to prefill workers) or any other scenario requiring intelligent KV cache-aware +routing decisions. +""" + +import asyncio +import logging +from typing import Optional + +import uvloop + +from dynamo.common.utils.token_ids import token_ids_to_list +from dynamo.llm import AisPerfConfig, KvRouter, KvRouterConfig +from dynamo.router.args import ( + DynamoRouterConfig, + build_ais_perf_config, + build_kv_router_config, +) +from dynamo.router.args import parse_args as parse_router_args +from dynamo.runtime import Client, DistributedRuntime, dynamo_worker +from dynamo.runtime.logging import configure_dynamo_logging + +configure_dynamo_logging() +logger = logging.getLogger(__name__) + + +class StandaloneRouterHandler: + """Handles routing requests to workers using KV-aware routing.""" + + def __init__( + self, + runtime: DistributedRuntime, + worker_endpoint_path: str, + block_size: int, + kv_router_config: KvRouterConfig, + ais_perf_config: Optional[AisPerfConfig], + ): + self.runtime = runtime + self.worker_endpoint_path = worker_endpoint_path + self.block_size = block_size + self.kv_router_config = kv_router_config + self.ais_perf_config = ais_perf_config + self.kv_router: Optional[KvRouter] = None + self.worker_client: Optional[Client] = None + + async def initialize(self): + """Initialize the KV router for workers.""" + try: + # Parse endpoint path (format: namespace.component.endpoint) + parts = self.worker_endpoint_path.split(".") + if len(parts) != 3: + raise ValueError( + f"Invalid endpoint path format: {self.worker_endpoint_path}. " + "Expected format: namespace.component.endpoint" + ) + namespace, component, endpoint = parts + + # Get worker endpoint + worker_endpoint = self.runtime.endpoint( + f"{namespace}.{component}.{endpoint}" + ) + self.worker_client = await worker_endpoint.client() + + self.kv_router = KvRouter( + endpoint=worker_endpoint, + block_size=self.block_size, + kv_router_config=self.kv_router_config, + ais_perf_config=self.ais_perf_config, + ) + + except Exception as e: + logger.error(f"Failed to initialize KvRouter: {e}") + raise + + async def generate(self, request): + """ + Generate tokens using the KV-aware router. + + Routes a PreprocessedRequest-shaped request to the best worker and streams back + its LLMEngineOutput responses unchanged, so fields this router does not inspect + (multimodal data, KV hints, agent context, ...) survive the hop. + """ + if self.kv_router is None: + logger.error("KvRouter not initialized - cannot process request") + raise RuntimeError("Router not initialized") + + preprocessed_request = dict(request) + preprocessed_request.setdefault("model", "unknown") + # Legacy callers send a top-level dp_rank instead of routing hints. + dp_rank = preprocessed_request.get("dp_rank") + if preprocessed_request.get("routing") is None and dp_rank is not None: + preprocessed_request["routing"] = {"dp_rank": dp_rank} + + async for worker_output in await self.kv_router.generate_from_request( + preprocessed_request + ): + yield worker_output + + async def best_worker_id( + self, token_ids, router_config_override=None, cache_namespace=None + ): + """ + Get the best worker ID for a given set of tokens without actually routing. + + This method returns the worker ID that would be selected based on KV cache + overlap, but does NOT actually route the request or update router states. + It's useful for debugging, monitoring, or implementing custom routing logic. + """ + if self.kv_router is None: + logger.error("KvRouter not initialized - cannot get best worker") + raise RuntimeError("Router not initialized") + + (worker_id, _dp_rank, _overlap_blocks) = await self.kv_router.best_worker( + token_ids, + router_config_override, + cache_namespace=cache_namespace, + ) + + yield worker_id + + async def get_overlap_scores(self, request): + """ + Get per-worker KV overlap by storage tier without routing the request. + + This endpoint returns matched blocks for each worker_id/dp_rank pair. + Shared-cache hits are request-global and are also reported per row as + blocks beyond that rank's device-local prefix. + """ + if self.kv_router is None: + logger.error("KvRouter not initialized - cannot get overlap scores") + raise RuntimeError("Router not initialized") + + scores = await self.kv_router.get_overlap_scores( + token_ids_to_list(request["token_ids"]), + request.get("router_config_override"), + request.get("block_mm_infos"), + request.get("lora_name"), + request.get("include_shared", True), + request.get("cache_namespace"), + ) + + yield scores + + +def parse_args(argv=None) -> DynamoRouterConfig: + """Parse router CLI arguments (compatibility shim delegating to args.parse_args).""" + return parse_router_args(argv) + + +@dynamo_worker() +async def worker(runtime: DistributedRuntime): + """Main worker function for the standalone router service.""" + + config = parse_args() + + logger.info("Starting Standalone Router Service") + logger.debug( + "Configuration: endpoint=%s, router_block_size=%s, " + "overlap_score_credit=%s, overlap_score_credit_decay=%s, " + "prefill_load_scale=%s, decode_active_request_weight=%s, " + "router_temperature=%s, use_kv_events=%s, router_replica_sync=%s, " + "router_track_active_blocks=%s, router_track_output_blocks=%s, " + "router_assume_kv_reuse=%s, router_track_prefill_tokens=%s, " + "router_ttl_secs=%s, router_approximate_cache_policy=%s", + config.endpoint, + config.router_block_size, + config.overlap_score_credit, + config.overlap_score_credit_decay, + config.prefill_load_scale, + config.decode_active_request_weight, + config.router_temperature, + config.use_kv_events, + config.router_replica_sync, + config.router_track_active_blocks, + config.router_track_output_blocks, + config.router_assume_kv_reuse, + config.router_track_prefill_tokens, + config.router_ttl_secs, + config.router_approximate_cache_policy, + ) + + kv_router_config = build_kv_router_config(config) + ais_perf_config = build_ais_perf_config(config) + + # Create handler + handler = StandaloneRouterHandler( + runtime, + config.endpoint, + config.router_block_size, + kv_router_config, + ais_perf_config, + ) + await handler.initialize() + + # Create endpoints + generate_endpoint = runtime.endpoint(f"{config.namespace}.router.generate") + best_worker_endpoint = runtime.endpoint(f"{config.namespace}.router.best_worker_id") + overlap_scores_endpoint = runtime.endpoint( + f"{config.namespace}.router.get_overlap_scores" + ) + + logger.debug("Starting to serve endpoints...") + + # Serve both endpoints concurrently + try: + await asyncio.gather( + generate_endpoint.serve_endpoint( + handler.generate, + graceful_shutdown=True, + metrics_labels=[("service", "router")], + ), + best_worker_endpoint.serve_endpoint( + handler.best_worker_id, + graceful_shutdown=True, + metrics_labels=[("service", "router")], + ), + overlap_scores_endpoint.serve_endpoint( + handler.get_overlap_scores, + graceful_shutdown=True, + metrics_labels=[("service", "router")], + ), + ) + except Exception as e: + logger.error(f"Failed to serve endpoint: {e}") + raise + finally: + logger.info("Standalone Router Service shutting down") + + +def main(): + """Entry point for the standalone router service.""" + uvloop.run(worker()) + + +if __name__ == "__main__": + main() diff --git a/components/src/dynamo/router/tests/test_router_jemalloc.py b/components/src/dynamo/router/tests/test_router_jemalloc.py new file mode 100644 index 000000000000..1a4a9a25c04c --- /dev/null +++ b/components/src/dynamo/router/tests/test_router_jemalloc.py @@ -0,0 +1,55 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Router 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 +import sys +from pathlib import Path +from types import ModuleType +from unittest.mock import Mock + +import pytest + +pytestmark = [pytest.mark.pre_merge, pytest.mark.unit, pytest.mark.gpu_0] + +ENTRYPOINT = Path(__file__).resolve().parents[1] / "__main__.py" + + +@pytest.fixture +def execve(monkeypatch): + for name in ("DYN_JEMALLOC", "DYN_FRONTEND_JEMALLOC", "LD_PRELOAD"): + monkeypatch.delenv(name, raising=False) + execve = Mock(side_effect=SystemExit) + monkeypatch.setattr( + ctypes.util, "find_library", Mock(return_value="libjemalloc.so.2") + ) + monkeypatch.setattr(os, "execve", execve) + return execve + + +def test_preloads_before_runtime_import(execve, monkeypatch): + monkeypatch.setenv("DYN_JEMALLOC", "1") + # Importing the runtime before the re-exec raises ImportError, not SystemExit. + monkeypatch.setitem(sys.modules, "dynamo.router.main", None) + with pytest.raises(SystemExit): + runpy.run_path(str(ENTRYPOINT), run_name="__main__") + _, _, env = execve.call_args.args + assert env["LD_PRELOAD"] == "libjemalloc.so.2" + + +@pytest.mark.parametrize("env_var", [None, "DYN_FRONTEND_JEMALLOC"]) +def test_runs_without_preload(execve, monkeypatch, env_var): + if env_var is not None: + monkeypatch.setenv(env_var, "1") + runtime = ModuleType("dynamo.router.main") + runtime.main = Mock() + monkeypatch.setitem(sys.modules, "dynamo.router.main", runtime) + runpy.run_path(str(ENTRYPOINT), run_name="__main__") + execve.assert_not_called() + runtime.main.assert_called_once_with() diff --git a/components/src/dynamo/router/tests/test_standalone_router.py b/components/src/dynamo/router/tests/test_standalone_router.py index b9aacaa6c626..0e1c1c721dac 100644 --- a/components/src/dynamo/router/tests/test_standalone_router.py +++ b/components/src/dynamo/router/tests/test_standalone_router.py @@ -65,7 +65,7 @@ def load_standalone_router_handler(): previous = {name: sys.modules.get(name) for name in stubs} sys.modules.update(stubs) try: - module_path = Path(__file__).parents[1] / "__main__.py" + module_path = Path(__file__).parents[1] / "main.py" spec = importlib.util.spec_from_file_location( "standalone_router_main", module_path ) diff --git a/docs/fern/pages/developer-guide/knowledge-base/modular-components/router/standalone-indexer.md b/docs/fern/pages/developer-guide/knowledge-base/modular-components/router/standalone-indexer.md index 264715c360ba..72a80b895f0c 100644 --- a/docs/fern/pages/developer-guide/knowledge-base/modular-components/router/standalone-indexer.md +++ b/docs/fern/pages/developer-guide/knowledge-base/modular-components/router/standalone-indexer.md @@ -142,6 +142,14 @@ For the standalone indexer, the gate runs before the HTTP listener is bound, so only be satisfied by initial `--workers` and workers recovered through `--peers`; `/register` cannot satisfy it. Leave the variable unset or set it to `0` to disable the startup wait. +### Host Memory Allocator + +Set `DYN_JEMALLOC=1` to preload jemalloc into the indexer process. The entry point adds the library +to `LD_PRELOAD` and re-executes itself once before it loads the indexer. If the library is not +found, the indexer prints a warning and starts with the default allocator. See +[Host memory allocator](../../../../reference/components/frontend-configuration.mdx#host-memory-allocator) +for the accepted values and for which container images ship `libjemalloc2`. + ## HTTP API ### `GET /health` — Liveness check diff --git a/docs/fern/pages/reference/components/frontend-configuration.mdx b/docs/fern/pages/reference/components/frontend-configuration.mdx index 01a6ca2a25c0..1a0196a24134 100644 --- a/docs/fern/pages/reference/components/frontend-configuration.mdx +++ b/docs/fern/pages/reference/components/frontend-configuration.mdx @@ -746,6 +746,22 @@ glibc allocator. volume with an initContainer or use a custom image. + + Preload jemalloc into the process. Accepts `1`, `true`, or `yes` (case-insensitive). At startup, + the entry point looks up the library with `ctypes.util.find_library("jemalloc")`, adds it to the + front of `LD_PRELOAD`, and re-executes itself once. Child processes inherit the preload. If + `LD_PRELOAD` already names a `libjemalloc` library, the process starts without re-executing. If + the library is not found or the re-execution fails, the process prints a warning and continues + with the default allocator. + + The lookup searches the dynamic linker cache and may not find a library copied into a shared + volume; set `LD_PRELOAD` to the copied path instead. + + The same variable applies to the standalone router (`python -m dynamo.router`) and the standalone + KV indexer (`python -m dynamo.indexer`). The frontend also accepts `DYN_FRONTEND_JEMALLOC` with + the same values and preloads jemalloc when either variable is enabled. + + Preload jemalloc for the frontend process. The library path depends on the container architecture: diff --git a/lib/bindings/python/src/dynamo/_jemalloc.py b/lib/bindings/python/src/dynamo/_jemalloc.py new file mode 100644 index 000000000000..db3323db890d --- /dev/null +++ b/lib/bindings/python/src/dynamo/_jemalloc.py @@ -0,0 +1,64 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Opt-in jemalloc preload for Dynamo Python entry points. + +Entry points call ``maybe_preload_jemalloc`` before importing the Dynamo runtime. +This module uses only the standard library so importing it does not load the +Rust extension. +""" + +import ctypes.util +import os +import sys + +JEMALLOC_ENV = "DYN_JEMALLOC" + + +def _enabled(name: str) -> bool: + # Match env_bool without importing Dynamo utilities before the re-exec. + return os.environ.get(name, "").lower() in ("1", "true", "yes") + + +def maybe_preload_jemalloc(alias: str | None = None) -> None: + """Re-exec the current process with jemalloc preloaded when requested. + + Set DYN_JEMALLOC, or the entry point's ``alias`` variable, 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. + """ + enabled_by = next( + (name for name in (JEMALLOC_ENV, alias) if name and _enabled(name)), None + ) + if enabled_by is None: + 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( + f"WARNING: {enabled_by} 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: {enabled_by} is enabled but re-exec with jemalloc failed: " + f"{exc}; continuing with the current allocator.", + file=sys.stderr, + ) diff --git a/lib/bindings/python/src/dynamo/indexer/__main__.py b/lib/bindings/python/src/dynamo/indexer/__main__.py index 504f64c74181..47b87c6bc6f9 100644 --- a/lib/bindings/python/src/dynamo/indexer/__main__.py +++ b/lib/bindings/python/src/dynamo/indexer/__main__.py @@ -1,7 +1,12 @@ # SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 -from dynamo.indexer.main import main +from dynamo._jemalloc import maybe_preload_jemalloc if __name__ == "__main__": + maybe_preload_jemalloc() + + # Import the runtime only after configuring the allocator. + from dynamo.indexer.main import main + raise SystemExit(main()) diff --git a/lib/bindings/python/tests/test_jemalloc_preload.py b/lib/bindings/python/tests/test_jemalloc_preload.py new file mode 100644 index 000000000000..4e3a7023ca8a --- /dev/null +++ b/lib/bindings/python/tests/test_jemalloc_preload.py @@ -0,0 +1,162 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +import ctypes.util +import os +import runpy +import sys +from pathlib import Path +from types import ModuleType +from unittest.mock import Mock + +import pytest + +from dynamo._jemalloc import maybe_preload_jemalloc + +pytestmark = [pytest.mark.pre_merge, pytest.mark.unit, pytest.mark.gpu_0] + +ALIAS = "DYN_FRONTEND_JEMALLOC" +INDEXER_ENTRYPOINT = ( + Path(__file__).resolve().parents[1] / "src" / "dynamo" / "indexer" / "__main__.py" +) + +# Environment, alias passed by the entry point, and the variable warnings name. +ENABLED = [ + pytest.param({"DYN_JEMALLOC": "1"}, None, "DYN_JEMALLOC", id="generic"), + pytest.param({"DYN_JEMALLOC": "true"}, ALIAS, "DYN_JEMALLOC", id="generic-alias"), + pytest.param({ALIAS: "YES"}, ALIAS, ALIAS, id="alias"), + pytest.param({"DYN_JEMALLOC": "0", ALIAS: "1"}, ALIAS, ALIAS, id="alias-only"), +] + + +@pytest.fixture +def startup(monkeypatch): + for name in ("DYN_JEMALLOC", ALIAS, "LD_PRELOAD"): + monkeypatch.delenv(name, raising=False) + find_library = Mock(return_value="libjemalloc.so.2") + execve = Mock(side_effect=SystemExit) + monkeypatch.setattr(ctypes.util, "find_library", find_library) + monkeypatch.setattr(os, "execve", execve) + return find_library, execve + + +def set_env(monkeypatch, env): + for name, value in env.items(): + monkeypatch.setenv(name, value) + + +@pytest.mark.parametrize("alias", [None, ALIAS]) +@pytest.mark.parametrize("value", [None, "", "0", "false", "no"]) +def test_disabled(startup, monkeypatch, alias, value): + if value is not None: + set_env(monkeypatch, {"DYN_JEMALLOC": value, ALIAS: value}) + find_library, execve = startup + maybe_preload_jemalloc(alias=alias) + find_library.assert_not_called() + execve.assert_not_called() + assert "LD_PRELOAD" not in os.environ + + +def test_alias_ignored_unless_passed(startup, monkeypatch): + monkeypatch.setenv(ALIAS, "1") + find_library, execve = startup + maybe_preload_jemalloc() + find_library.assert_not_called() + execve.assert_not_called() + + +@pytest.mark.parametrize("env,alias,enabled_by", ENABLED) +@pytest.mark.parametrize("separator", [":", " "]) +def test_already_preloaded(startup, monkeypatch, env, alias, enabled_by, separator): + set_env(monkeypatch, env) + preload = f"libother.so{separator}/usr/lib/libjemalloc.so.2" + monkeypatch.setenv("LD_PRELOAD", preload) + find_library, execve = startup + maybe_preload_jemalloc(alias=alias) + find_library.assert_not_called() + execve.assert_not_called() + assert os.environ["LD_PRELOAD"] == preload + + +@pytest.mark.parametrize("env,alias,enabled_by", ENABLED) +def test_missing_library(startup, monkeypatch, capsys, env, alias, enabled_by): + set_env(monkeypatch, env) + find_library, execve = startup + find_library.return_value = None + maybe_preload_jemalloc(alias=alias) + find_library.assert_called_once_with("jemalloc") + execve.assert_not_called() + err = capsys.readouterr().err + assert f"{enabled_by} is enabled but libjemalloc was not found" in err + assert "LD_PRELOAD" not in os.environ + + +@pytest.mark.parametrize("env,alias,enabled_by", ENABLED) +@pytest.mark.parametrize( + "existing", ["", "libother.so:libanother.so", "/opt/jemalloc-tools/libheaptrace.so"] +) +def test_restart(startup, monkeypatch, env, alias, enabled_by, existing): + set_env(monkeypatch, env) + monkeypatch.setenv("LD_PRELOAD", existing) + argv = [sys.executable, "-u", "-X", "faulthandler", "-m", "dynamo.router"] + argv += ["--endpoint", "dynamo.prefill.generate"] + monkeypatch.setattr(sys, "orig_argv", argv) + find_library, execve = 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 "") + ) + with pytest.raises(SystemExit): + maybe_preload_jemalloc(alias=alias) + find_library.assert_called_once_with("jemalloc") + execve.assert_called_once_with(sys.executable, argv, expected_env) + assert os.environ["LD_PRELOAD"] == existing + + +@pytest.mark.parametrize("env,alias,enabled_by", ENABLED) +@pytest.mark.parametrize("existing", [None, "", "libother.so"]) +def test_exec_failure(startup, monkeypatch, capsys, env, alias, enabled_by, existing): + set_env(monkeypatch, env) + if existing is not None: + monkeypatch.setenv("LD_PRELOAD", existing) + _, execve = startup + execve.side_effect = OSError("exec blocked") + maybe_preload_jemalloc(alias=alias) + assert os.environ.get("LD_PRELOAD") == existing + err = capsys.readouterr().err + assert f"{enabled_by} is enabled but re-exec with jemalloc failed" in err + assert "exec blocked" in err + + +def test_indexer_preloads_before_runtime_import(startup, monkeypatch): + monkeypatch.setenv("DYN_JEMALLOC", "1") + # Importing the runtime before the re-exec raises ImportError, not SystemExit. + monkeypatch.setitem(sys.modules, "dynamo.indexer.main", None) + _, execve = startup + with pytest.raises(SystemExit): + runpy.run_path(str(INDEXER_ENTRYPOINT), run_name="__main__") + _, _, env = execve.call_args.args + assert env["LD_PRELOAD"] == "libjemalloc.so.2" + + +@pytest.mark.parametrize("env", [{}, {ALIAS: "1"}]) +def test_indexer_runs_without_preload(startup, monkeypatch, env): + set_env(monkeypatch, env) + runtime = ModuleType("dynamo.indexer.main") + runtime.main = Mock(return_value=0) + monkeypatch.setitem(sys.modules, "dynamo.indexer.main", runtime) + _, execve = startup + with pytest.raises(SystemExit) as exit_info: + runpy.run_path(str(INDEXER_ENTRYPOINT), run_name="__main__") + assert exit_info.value.code == 0 + execve.assert_not_called() + runtime.main.assert_called_once_with() From c6a20fa61ca1cb6d7c5034b33460a77a74876fad Mon Sep 17 00:00:00 2001 From: PeaBrane Date: Wed, 7 Oct 2026 22:32:02 -0700 Subject: [PATCH 2/2] test(bindings): trim enabled-source matrix from jemalloc preload tests The already-preloaded and restart paths do not depend on which variable enabled the helper, so run them under DYN_JEMALLOC=1 only. Alias and selection coverage stays in test_alias_ignored_unless_passed and the ENABLED matrix in test_missing_library and test_exec_failure. Signed-off-by: PeaBrane --- lib/bindings/python/tests/test_jemalloc_preload.py | 14 ++++++-------- 1 file changed, 6 insertions(+), 8 deletions(-) diff --git a/lib/bindings/python/tests/test_jemalloc_preload.py b/lib/bindings/python/tests/test_jemalloc_preload.py index 4e3a7023ca8a..506559e4e0ca 100644 --- a/lib/bindings/python/tests/test_jemalloc_preload.py +++ b/lib/bindings/python/tests/test_jemalloc_preload.py @@ -65,14 +65,13 @@ def test_alias_ignored_unless_passed(startup, monkeypatch): execve.assert_not_called() -@pytest.mark.parametrize("env,alias,enabled_by", ENABLED) @pytest.mark.parametrize("separator", [":", " "]) -def test_already_preloaded(startup, monkeypatch, env, alias, enabled_by, separator): - set_env(monkeypatch, env) +def test_already_preloaded(startup, monkeypatch, separator): + monkeypatch.setenv("DYN_JEMALLOC", "1") preload = f"libother.so{separator}/usr/lib/libjemalloc.so.2" monkeypatch.setenv("LD_PRELOAD", preload) find_library, execve = startup - maybe_preload_jemalloc(alias=alias) + maybe_preload_jemalloc() find_library.assert_not_called() execve.assert_not_called() assert os.environ["LD_PRELOAD"] == preload @@ -91,12 +90,11 @@ def test_missing_library(startup, monkeypatch, capsys, env, alias, enabled_by): assert "LD_PRELOAD" not in os.environ -@pytest.mark.parametrize("env,alias,enabled_by", ENABLED) @pytest.mark.parametrize( "existing", ["", "libother.so:libanother.so", "/opt/jemalloc-tools/libheaptrace.so"] ) -def test_restart(startup, monkeypatch, env, alias, enabled_by, existing): - set_env(monkeypatch, env) +def test_restart(startup, monkeypatch, existing): + monkeypatch.setenv("DYN_JEMALLOC", "1") monkeypatch.setenv("LD_PRELOAD", existing) argv = [sys.executable, "-u", "-X", "faulthandler", "-m", "dynamo.router"] argv += ["--endpoint", "dynamo.prefill.generate"] @@ -116,7 +114,7 @@ def replace_process(*args): os.environ, LD_PRELOAD="libjemalloc.so.2" + (f":{existing}" if existing else "") ) with pytest.raises(SystemExit): - maybe_preload_jemalloc(alias=alias) + maybe_preload_jemalloc() find_library.assert_called_once_with("jemalloc") execve.assert_called_once_with(sys.executable, argv, expected_env) assert os.environ["LD_PRELOAD"] == existing