diff --git a/components/src/dynamo/common/configuration/groups/runtime_args.py b/components/src/dynamo/common/configuration/groups/runtime_args.py index ab646b1f8737..b1887c67cfe7 100644 --- a/components/src/dynamo/common/configuration/groups/runtime_args.py +++ b/components/src/dynamo/common/configuration/groups/runtime_args.py @@ -4,6 +4,8 @@ """Dynamo runtime configuration ArgGroup.""" import argparse +import logging +import os from typing import List, Optional from dynamo._core import get_reasoning_parser_names, get_tool_parser_names @@ -13,6 +15,10 @@ from dynamo.common.utils.namespace import get_worker_namespace from dynamo.common.utils.output_modalities import OutputModality +logger = logging.getLogger(__name__) +_FPM_TRACE_VALUES = {"1", "0", "true", "false", "on", "off", "yes", "no"} +_fpm_trace_invalid_warning_emitted = False + class DynamoRuntimeConfig(ConfigBase): """Configuration for Dynamo runtime (common across all backends).""" @@ -22,6 +28,7 @@ class DynamoRuntimeConfig(ConfigBase): discovery_backend: str request_plane: str event_plane: Optional[str] = None + fpm_trace: bool = False connector: list[str] enable_local_indexer: bool durable_kv_events: bool @@ -52,6 +59,29 @@ class DynamoRuntimeConfig(ConfigBase): def validate(self) -> None: self.namespace = get_worker_namespace(self.namespace) + # The Rust FPM sink reads this setting from the process environment. + # Canonicalize the resolved CLI/env value before the runtime or backend + # child processes are created so --fpm-trace and --no-fpm-trace apply to + # both the Python instrumentation and the Rust persistence layer. + if self.fpm_trace or "DYN_FPM_TRACE" in os.environ: + raw_fpm_trace = os.environ.get("DYN_FPM_TRACE") + if ( + raw_fpm_trace is not None + and raw_fpm_trace.strip().lower() not in _FPM_TRACE_VALUES + and not self.fpm_trace + and "DYN_FORWARDPASS_METRIC_PORT" not in os.environ + ): + global _fpm_trace_invalid_warning_emitted + if not _fpm_trace_invalid_warning_emitted: + _fpm_trace_invalid_warning_emitted = True + logger.warning( + "Invalid DYN_FPM_TRACE value %r; expected one of 1/0, " + "true/false, on/off, or yes/no. FPM tracing is disabled " + "for this worker.", + raw_fpm_trace, + ) + os.environ["DYN_FPM_TRACE"] = "1" if self.fpm_trace else "0" + # TODO get a better way for spot fixes like this. self.enable_local_indexer = not self.durable_kv_events self._validate_output_modalities() @@ -126,6 +156,13 @@ def add_arguments(self, parser: argparse.ArgumentParser) -> None: "all discovery backends. Set to 'nats' to use a NATS-based event plane.", choices=["nats", "zmq"], ) + add_negatable_bool_argument( + g, + flag_name="--fpm-trace", + env_var="DYN_FPM_TRACE", + default=False, + help="Persist backend forward-pass metrics to rotating gzip JSONL trace files. Also enables the backend FPM instrumentation required to produce those records.", + ) add_argument( g, flag_name="--connector", diff --git a/components/src/dynamo/common/configuration/utils.py b/components/src/dynamo/common/configuration/utils.py index 3e50b8e61d42..ca259f60cde1 100644 --- a/components/src/dynamo/common/configuration/utils.py +++ b/components/src/dynamo/common/configuration/utils.py @@ -42,7 +42,7 @@ def env_or_default( target_type = value_type if value_type is not None else type(default) if target_type is bool: - return value.lower() in ("true", "1", "yes", "on") # type: ignore + return value.strip().lower() in ("true", "1", "yes", "on") # type: ignore if target_type is int: return int(value) # type: ignore if target_type is float: diff --git a/components/src/dynamo/common/tests/configuration/test_runtime_args.py b/components/src/dynamo/common/tests/configuration/test_runtime_args.py new file mode 100644 index 000000000000..973e33594567 --- /dev/null +++ b/components/src/dynamo/common/tests/configuration/test_runtime_args.py @@ -0,0 +1,113 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Tests for shared Dynamo runtime arguments.""" + +import argparse +import logging +import os + +import pytest + +import dynamo.common.configuration.groups.runtime_args as runtime_args +from dynamo.common.configuration.groups.runtime_args import ( + DynamoRuntimeArgGroup, + DynamoRuntimeConfig, +) + +pytestmark = [ + pytest.mark.unit, + pytest.mark.gpu_0, + pytest.mark.pre_merge, +] + + +def _parse_runtime_args(argv: list[str]) -> tuple[DynamoRuntimeConfig, str]: + parser = argparse.ArgumentParser() + DynamoRuntimeArgGroup().add_arguments(parser) + args = parser.parse_args(argv) + config = DynamoRuntimeConfig.from_cli_args(args) + config.validate() + return config, parser.format_help() + + +def test_fpm_trace_defaults_disabled(monkeypatch): + monkeypatch.delenv("DYN_FPM_TRACE", raising=False) + + config, _ = _parse_runtime_args([]) + + assert config.fpm_trace is False + assert "DYN_FPM_TRACE" not in os.environ + + +def test_fpm_trace_env_enables_and_is_canonicalized(monkeypatch): + monkeypatch.setenv("DYN_FPM_TRACE", "on") + + config, _ = _parse_runtime_args([]) + + assert config.fpm_trace is True + assert os.environ["DYN_FPM_TRACE"] == "1" + + +def test_fpm_trace_env_is_trimmed(monkeypatch): + monkeypatch.setenv("DYN_FPM_TRACE", " true ") + + config, _ = _parse_runtime_args([]) + + assert config.fpm_trace is True + assert os.environ["DYN_FPM_TRACE"] == "1" + + +def test_invalid_fpm_trace_warns_once_and_is_disabled(monkeypatch, caplog): + monkeypatch.setenv("DYN_FPM_TRACE", "sometimes") + monkeypatch.setattr(runtime_args, "_fpm_trace_invalid_warning_emitted", False) + + with caplog.at_level(logging.WARNING, logger=runtime_args.__name__): + config, _ = _parse_runtime_args([]) + monkeypatch.setenv("DYN_FPM_TRACE", "still-invalid") + _parse_runtime_args([]) + + assert config.fpm_trace is False + assert os.environ["DYN_FPM_TRACE"] == "0" + assert caplog.text.count("Invalid DYN_FPM_TRACE value") == 1 + + +def test_explicit_fpm_port_preserves_precedence_over_invalid_trace(monkeypatch, caplog): + monkeypatch.setenv("DYN_FORWARDPASS_METRIC_PORT", "23456") + monkeypatch.setenv("DYN_FPM_TRACE", "sometimes") + monkeypatch.setattr(runtime_args, "_fpm_trace_invalid_warning_emitted", False) + + with caplog.at_level(logging.WARNING, logger=runtime_args.__name__): + config, _ = _parse_runtime_args([]) + + assert config.fpm_trace is False + assert os.environ["DYN_FPM_TRACE"] == "0" + assert "Invalid DYN_FPM_TRACE value" not in caplog.text + + +def test_fpm_trace_cli_enables_and_is_exported(monkeypatch): + monkeypatch.delenv("DYN_FPM_TRACE", raising=False) + + config, _ = _parse_runtime_args(["--fpm-trace"]) + + assert config.fpm_trace is True + assert os.environ["DYN_FPM_TRACE"] == "1" + + +def test_no_fpm_trace_cli_overrides_enabled_env(monkeypatch): + monkeypatch.setenv("DYN_FPM_TRACE", "true") + + config, _ = _parse_runtime_args(["--no-fpm-trace"]) + + assert config.fpm_trace is False + assert os.environ["DYN_FPM_TRACE"] == "0" + + +def test_fpm_trace_help_lists_flag_and_env(monkeypatch): + monkeypatch.delenv("DYN_FPM_TRACE", raising=False) + + _, help_text = _parse_runtime_args([]) + + assert "--fpm-trace" in help_text + assert "--no-fpm-trace" in help_text + assert "DYN_FPM_TRACE" in help_text diff --git a/components/src/dynamo/common/tests/configuration/test_utils.py b/components/src/dynamo/common/tests/configuration/test_utils.py index bd2e6020848d..7952966a4748 100644 --- a/components/src/dynamo/common/tests/configuration/test_utils.py +++ b/components/src/dynamo/common/tests/configuration/test_utils.py @@ -39,7 +39,7 @@ def test_returns_env_when_set(self, monkeypatch): def test_bool_conversion_true(self, monkeypatch): """Test bool conversion for true values.""" - test_cases = ["true", "True", "1", "yes", "YES", "on", "ON"] + test_cases = ["true", "True", "1", "yes", "YES", "on", "ON", " true "] for value in test_cases: monkeypatch.setenv("TEST_BOOL", value) @@ -48,7 +48,7 @@ def test_bool_conversion_true(self, monkeypatch): def test_bool_conversion_false(self, monkeypatch): """Test bool conversion for false values.""" - test_cases = ["false", "False", "0", "no", "NO", "off", "OFF"] + test_cases = ["false", "False", "0", "no", "NO", "off", "OFF", " off "] for value in test_cases: monkeypatch.setenv("TEST_BOOL", value) diff --git a/components/src/dynamo/common/utils/tests/test_env.py b/components/src/dynamo/common/utils/tests/test_env.py index de8dbc973569..6dc66b29fa70 100644 --- a/components/src/dynamo/common/utils/tests/test_env.py +++ b/components/src/dynamo/common/utils/tests/test_env.py @@ -30,7 +30,7 @@ def test_truthy_values(self, monkeypatch, value): monkeypatch.setenv("FOO", value) assert env_bool("FOO") is True - @pytest.mark.parametrize("value", ["false", "0", "no", "anything"]) + @pytest.mark.parametrize("value", ["false", "0", "no", "on", "off", "anything"]) def test_falsy_values(self, monkeypatch, value): monkeypatch.setenv("FOO", value) assert env_bool("FOO") is False diff --git a/components/src/dynamo/sglang/args.py b/components/src/dynamo/sglang/args.py index d3296ad2b6a9..8daca61cd9c2 100644 --- a/components/src/dynamo/sglang/args.py +++ b/components/src/dynamo/sglang/args.py @@ -66,6 +66,48 @@ def _set_serving_strategy(self): return DisaggregationMode.AGGREGATED +def _unsupported_fpm_trace_role(dynamo_config: DynamoConfig) -> Optional[str]: + """Return the worker role when the selected path does not create an FPM relay.""" + if is_snapshot_enabled(): + return "snapshot" + if dynamo_config.embedding_worker: + return "embedding" + if ( + dynamo_config.multimodal_encode_worker + or dynamo_config.multimodal_worker + or dynamo_config.dedicated_mm_encoder + ): + return "dedicated multimodal" + if dynamo_config.image_diffusion_worker: + return "image diffusion" + if dynamo_config.video_generation_worker: + return "video generation" + return None + + +def _forward_pass_metrics_source( + dynamo_config: DynamoConfig, *, fpm_trace_relay_supported: bool = True +) -> Optional[str]: + """Resolve the FPM opt-in source while preserving the legacy port switch.""" + if os.environ.get("DYN_FORWARDPASS_METRIC_PORT"): + return "DYN_FORWARDPASS_METRIC_PORT" + if not dynamo_config.fpm_trace: + return None + + unsupported_role = _unsupported_fpm_trace_role(dynamo_config) + if unsupported_role is None and not fpm_trace_relay_supported: + unsupported_role = "unified backend" + if unsupported_role is None: + return "--fpm-trace/DYN_FPM_TRACE" + + logging.warning( + "--fpm-trace/DYN_FPM_TRACE is enabled, but SGLang %s workers do not create a Dynamo " + "FPM relay. Trace-based FPM activation is disabled for this worker.", + unsupported_role, + ) + return None + + def use_modelexpress_remote_instance(args: Any) -> bool: return ( getattr(args, "load_format", None) == "remote_instance" @@ -264,12 +306,16 @@ def _dump_disagg_config_section(disagg_config: dict[str, Any]) -> str: return temp_path -async def parse_args(args: list[str]) -> Config: +async def parse_args( + args: list[str], *, fpm_trace_relay_supported: bool = True +) -> Config: """Parse CLI arguments and return combined configuration. Download the model if necessary. Args: args: Command-line argument strings. + fpm_trace_relay_supported: Whether this entry point constructs the + Dynamo relay required for trace-based FPM activation. Returns: Config object with server_args and dynamo_args. @@ -526,11 +572,13 @@ async def parse_args(args: list[str]) -> Config: ) # Enable forward pass metrics from dynamo env var if configured - if os.environ.get("DYN_FORWARDPASS_METRIC_PORT") and not getattr( - server_args, "enable_forward_pass_metrics", False - ): + fpm_source = _forward_pass_metrics_source( + dynamo_config, + fpm_trace_relay_supported=fpm_trace_relay_supported, + ) + if fpm_source and not getattr(server_args, "enable_forward_pass_metrics", False): server_args.enable_forward_pass_metrics = True - logging.info("Enabled forward_pass_metrics from DYN_FORWARDPASS_METRIC_PORT") + logging.info("Enabled forward_pass_metrics from %s", fpm_source) # Auto-detect diffusion worker mode if dllm_algorithm diffusion_worker = server_args.dllm_algorithm is not None diff --git a/components/src/dynamo/sglang/llm_engine.py b/components/src/dynamo/sglang/llm_engine.py index a4d2e48e9565..0be47b6bf07b 100644 --- a/components/src/dynamo/sglang/llm_engine.py +++ b/components/src/dynamo/sglang/llm_engine.py @@ -149,7 +149,10 @@ def __init__(self, server_args, dynamo_args, serving_mode: DisaggregationMode): async def from_args( cls, argv: list[str] | None = None ) -> tuple[SglangLLMEngine, WorkerConfig]: - config = await parse_args(argv if argv is not None else sys.argv[1:]) + config = await parse_args( + argv if argv is not None else sys.argv[1:], + fpm_trace_relay_supported=False, + ) server_args = config.server_args dynamo_args = config.dynamo_args diff --git a/components/src/dynamo/sglang/tests/test_sglang_unit.py b/components/src/dynamo/sglang/tests/test_sglang_unit.py index 6d49e36f304d..f271db417594 100644 --- a/components/src/dynamo/sglang/tests/test_sglang_unit.py +++ b/components/src/dynamo/sglang/tests/test_sglang_unit.py @@ -3,6 +3,8 @@ """Unit tests for SGLang backend components.""" +import logging +import os import re import sys from pathlib import Path @@ -13,12 +15,15 @@ from sglang.srt.disaggregation.utils import FAKE_BOOTSTRAP_HOST import dynamo.sglang._compat as sglang_compat -from dynamo.common.constants import EmbeddingTransferMode +import dynamo.sglang.llm_engine as sglang_llm_engine +from dynamo.common.constants import DisaggregationMode, EmbeddingTransferMode +from dynamo.common.snapshot.constants import SNAPSHOT_CONTROL_DIR_ENV from dynamo.sglang._compat import ( ensure_sglang_top_level_exports, filter_supported_async_generate_kwargs, ) from dynamo.sglang.args import ( + _forward_pass_metrics_source, _normalize_multimodal_disaggregation_args, parse_args, should_fetch_model, @@ -72,6 +77,7 @@ def _make_sglang_config(**overrides): config.enable_rl = False config.frontend_decoding = False config.sglang_trace_level = 2 + config.fpm_trace = False config.disagg_config = None config.disagg_config_key = None for key, value in overrides.items(): @@ -467,13 +473,172 @@ def test_dedicated_mm_encoder_requires_enable_multimodal(): @pytest.mark.asyncio async def test_forward_pass_metrics_enabled_from_env(monkeypatch, mock_sglang_cli): """Dynamo should enable FPM when DYN_FORWARDPASS_METRIC_PORT is set.""" - monkeypatch.setenv("DYN_FORWARDPASS_METRIC_PORT", "1") + monkeypatch.setenv("DYN_FORWARDPASS_METRIC_PORT", "23456") mock_sglang_cli("--model", "Qwen/Qwen3-0.6B") config = await parse_args(sys.argv[1:]) assert config.server_args.enable_forward_pass_metrics is True +@pytest.mark.asyncio +async def test_explicit_fpm_port_takes_precedence_over_trace( + monkeypatch, mock_sglang_cli, caplog +): + """The legacy explicit port remains authoritative even if trace is invalid.""" + monkeypatch.setenv("DYN_FORWARDPASS_METRIC_PORT", "23456") + monkeypatch.setenv("DYN_FPM_TRACE", "sometimes") + mock_sglang_cli("--model", "Qwen/Qwen3-0.6B") + + with caplog.at_level(logging.INFO): + config = await parse_args(sys.argv[1:]) + + assert config.server_args.enable_forward_pass_metrics is True + assert ( + "Enabled forward_pass_metrics from DYN_FORWARDPASS_METRIC_PORT" in caplog.text + ) + assert "Invalid DYN_FPM_TRACE value" not in caplog.text + + +@pytest.mark.asyncio +async def test_forward_pass_metrics_enabled_from_trace(monkeypatch, mock_sglang_cli): + """DYN_FPM_TRACE should enable SGLang's existing FPM publisher.""" + monkeypatch.delenv("DYN_FORWARDPASS_METRIC_PORT", raising=False) + monkeypatch.setenv("DYN_FPM_TRACE", "on") + mock_sglang_cli("--model", "Qwen/Qwen3-0.6B") + + config = await parse_args(sys.argv[1:]) + assert config.server_args.enable_forward_pass_metrics is True + + +@pytest.mark.asyncio +async def test_forward_pass_metrics_enabled_from_cli_flag(monkeypatch, mock_sglang_cli): + """The shared CLI flag should enable both Python and Rust trace handling.""" + monkeypatch.delenv("DYN_FORWARDPASS_METRIC_PORT", raising=False) + monkeypatch.delenv("DYN_FPM_TRACE", raising=False) + mock_sglang_cli("--fpm-trace", "--model", "Qwen/Qwen3-0.6B") + + config = await parse_args(sys.argv[1:]) + + assert config.dynamo_args.fpm_trace is True + assert config.server_args.enable_forward_pass_metrics is True + assert os.environ["DYN_FPM_TRACE"] == "1" + + +@pytest.mark.asyncio +async def test_false_fpm_trace_does_not_enable_metrics(monkeypatch, mock_sglang_cli): + monkeypatch.delenv("DYN_FORWARDPASS_METRIC_PORT", raising=False) + monkeypatch.setenv("DYN_FPM_TRACE", "off") + mock_sglang_cli("--model", "Qwen/Qwen3-0.6B") + + config = await parse_args(sys.argv[1:]) + assert not config.server_args.enable_forward_pass_metrics + + +@pytest.mark.asyncio +async def test_invalid_fpm_trace_is_disabled_by_arg_parser( + monkeypatch, mock_sglang_cli +): + monkeypatch.delenv("DYN_FORWARDPASS_METRIC_PORT", raising=False) + monkeypatch.setenv("DYN_FPM_TRACE", "sometimes") + mock_sglang_cli("--model", "Qwen/Qwen3-0.6B") + + config = await parse_args(sys.argv[1:]) + + assert not config.server_args.enable_forward_pass_metrics + assert os.environ["DYN_FPM_TRACE"] == "0" + + +@pytest.mark.parametrize( + ("overrides", "role", "fpm_trace_relay_supported"), + [ + ({}, "unified backend", False), + ({"embedding_worker": True}, "embedding", True), + ({"multimodal_encode_worker": True}, "dedicated multimodal", True), + ({"multimodal_worker": True}, "dedicated multimodal", True), + ({"image_diffusion_worker": True}, "image diffusion", True), + ({"video_generation_worker": True}, "video generation", True), + ], +) +def test_trace_does_not_activate_fpm_without_relay( + monkeypatch, caplog, overrides, role, fpm_trace_relay_supported +): + monkeypatch.delenv("DYN_FORWARDPASS_METRIC_PORT", raising=False) + dynamo_config = _make_sglang_config(fpm_trace=True, **overrides) + + with caplog.at_level(logging.WARNING): + source = _forward_pass_metrics_source( + dynamo_config, + fpm_trace_relay_supported=fpm_trace_relay_supported, + ) + + assert source is None + assert f"SGLang {role} workers do not create a Dynamo FPM relay" in caplog.text + + +def test_explicit_port_preserves_legacy_activation_without_relay(monkeypatch, caplog): + monkeypatch.setenv("DYN_FORWARDPASS_METRIC_PORT", "23456") + dynamo_config = _make_sglang_config(embedding_worker=True, fpm_trace=True) + + with caplog.at_level(logging.WARNING): + source = _forward_pass_metrics_source( + dynamo_config, + fpm_trace_relay_supported=False, + ) + + assert source == "DYN_FORWARDPASS_METRIC_PORT" + assert "do not create a Dynamo FPM relay" not in caplog.text + + +def test_trace_does_not_activate_fpm_during_snapshot_startup( + monkeypatch, caplog, tmp_path +): + monkeypatch.delenv("DYN_FORWARDPASS_METRIC_PORT", raising=False) + monkeypatch.setenv(SNAPSHOT_CONTROL_DIR_ENV, str(tmp_path)) + + with caplog.at_level(logging.WARNING): + source = _forward_pass_metrics_source(_make_sglang_config(fpm_trace=True)) + + assert source is None + assert "SGLang snapshot workers do not create a Dynamo FPM relay" in caplog.text + + +@pytest.mark.asyncio +async def test_unified_from_args_marks_fpm_relay_unsupported(monkeypatch): + server_args = SimpleNamespace( + skip_tokenizer_init=True, + model_path="Qwen/Qwen3-0.6B", + served_model_name="Qwen/Qwen3-0.6B", + ) + dynamo_args = SimpleNamespace(use_sglang_tokenizer=False) + config = SimpleNamespace( + server_args=server_args, + dynamo_args=dynamo_args, + serving_mode=DisaggregationMode.AGGREGATED, + ) + worker_config = object() + parse_options = {} + + async def fake_parse_args(argv, *, fpm_trace_relay_supported): + parse_options["fpm_trace_relay_supported"] = fpm_trace_relay_supported + return config + + monkeypatch.delenv("DYN_ENABLE_TEST_LOGITS_PROCESSOR", raising=False) + monkeypatch.setattr(sglang_llm_engine, "parse_args", fake_parse_args) + monkeypatch.setattr( + sglang_llm_engine.WorkerConfig, + "from_runtime_config", + lambda *args, **kwargs: worker_config, + ) + + engine, result_worker_config = await sglang_llm_engine.SglangLLMEngine.from_args( + ["--model-path", "Qwen/Qwen3-0.6B"] + ) + + assert engine.server_args is server_args + assert result_worker_config is worker_config + assert parse_options["fpm_trace_relay_supported"] is False + + @pytest.mark.asyncio async def test_obsolete_dyn_endpoint_types_flag_is_supported(mock_sglang_cli): """Obsolete --dyn-endpoint-types alias should map to endpoint_types.""" diff --git a/components/src/dynamo/vllm/args.py b/components/src/dynamo/vllm/args.py index 7dc58e87632c..8dac4fc513a8 100644 --- a/components/src/dynamo/vllm/args.py +++ b/components/src/dynamo/vllm/args.py @@ -63,11 +63,15 @@ def _preprocess_for_encode_config(config: Config) -> Dict[str, Any]: return config.__dict__ -def parse_args(argv: list[str] | None = None) -> Config: +def parse_args( + argv: list[str] | None = None, *, fpm_trace_relay_supported: bool = True +) -> Config: """Parse command-line arguments for the vLLM backend. Args: argv: Command-line arguments. ``None`` means ``sys.argv[1:]``. + fpm_trace_relay_supported: Whether this entry point constructs the + Dynamo relay required for trace-based FPM activation. Returns: Config: Parsed configuration object. @@ -114,7 +118,11 @@ def parse_args(argv: list[str] | None = None) -> Config: cross_validate_config(dynamo_config, engine_config) update_dynamo_config_with_engine(dynamo_config, engine_config) - update_engine_config_with_dynamo(dynamo_config, engine_config) + update_engine_config_with_dynamo( + dynamo_config, + engine_config, + fpm_trace_relay_supported=fpm_trace_relay_supported, + ) dynamo_config.engine_args = engine_config return dynamo_config @@ -226,8 +234,45 @@ def update_dynamo_config_with_engine( dynamo_config.connector = [] # type: ignore[assignment] +def _unsupported_fpm_trace_role(dynamo_config: Config) -> Optional[str]: + """Return the worker role when trace-based FPM activation is unsupported.""" + if dynamo_config.embedding_worker: + return "embedding" + if dynamo_config.headless: + return "headless" + if dynamo_config.disaggregation_mode == DisaggregationMode.ENCODE: + return "multimodal encode" + return None + + +def _forward_pass_metrics_enabled( + dynamo_config: Config, *, fpm_trace_relay_supported: bool = True +) -> bool: + """Resolve FPM activation without changing the legacy explicit-port path.""" + if envs.is_set("DYN_FORWARDPASS_METRIC_PORT"): + return True + if not dynamo_config.fpm_trace: + return False + + unsupported_role = _unsupported_fpm_trace_role(dynamo_config) + if unsupported_role is None and not fpm_trace_relay_supported: + unsupported_role = "unified backend" + if unsupported_role is None: + return True + + logger.warning( + "--fpm-trace/DYN_FPM_TRACE is enabled, but vLLM %s workers do not create a Dynamo " + "FPM relay. Trace-based FPM activation is disabled for this worker.", + unsupported_role, + ) + return False + + def update_engine_config_with_dynamo( - dynamo_config: Config, engine_config: AsyncEngineArgs + dynamo_config: Config, + engine_config: AsyncEngineArgs, + *, + fpm_trace_relay_supported: bool = True, ) -> None: """Update engine config based on Dynamo config.""" if engine_config.enable_prefix_caching is None: @@ -267,7 +312,11 @@ def update_engine_config_with_dynamo( f"(use_kv_events={dynamo_config.use_kv_events})" ) - if envs.is_set("DYN_FORWARDPASS_METRIC_PORT"): + fpm_enabled = _forward_pass_metrics_enabled( + dynamo_config, + fpm_trace_relay_supported=fpm_trace_relay_supported, + ) + if fpm_enabled: existing_cls = getattr(engine_config, "scheduler_cls", None) if existing_cls is None: defaults[ @@ -278,8 +327,13 @@ def update_engine_config_with_dynamo( f"(port={envs.DYN_FORWARDPASS_METRIC_PORT})" ) else: + fpm_source = ( + "DYN_FORWARDPASS_METRIC_PORT is set" + if envs.is_set("DYN_FORWARDPASS_METRIC_PORT") + else "--fpm-trace/DYN_FPM_TRACE is enabled" + ) logger.warning( - f"DYN_FORWARDPASS_METRIC_PORT is set but scheduler_cls " + f"{fpm_source} but scheduler_cls " f"is already '{existing_cls}'. InstrumentedScheduler will NOT " f"be injected. To use forward pass metrics, either remove " f"--scheduler-cls or subclass InstrumentedScheduler." @@ -292,7 +346,7 @@ def update_engine_config_with_dynamo( "Benchmark data will be collected but not served via endpoint." ) existing_cls = getattr(engine_config, "scheduler_cls", None) - if existing_cls is None and not envs.is_set("DYN_FORWARDPASS_METRIC_PORT"): + if existing_cls is None and not fpm_enabled: defaults[ "scheduler_cls" ] = "dynamo.vllm.instrumented_scheduler.InstrumentedScheduler" diff --git a/components/src/dynamo/vllm/envs.py b/components/src/dynamo/vllm/envs.py index f164f5cef80a..3398186f83aa 100644 --- a/components/src/dynamo/vllm/envs.py +++ b/components/src/dynamo/vllm/envs.py @@ -17,9 +17,10 @@ # Port range constants REGISTERED_PORT_MIN = 1024 REGISTERED_PORT_MAX = 49151 +DEFAULT_FORWARDPASS_METRIC_PORT = 20380 if TYPE_CHECKING: - DYN_FORWARDPASS_METRIC_PORT: int = 20380 + DYN_FORWARDPASS_METRIC_PORT: int = DEFAULT_FORWARDPASS_METRIC_PORT def _resolve_port(env_var: str, default_port: int) -> int: @@ -59,7 +60,7 @@ def _resolve_port(env_var: str, default_port: int) -> int: # Environment variables configuration environment_variables: dict[str, Callable[[], Any]] = { "DYN_FORWARDPASS_METRIC_PORT": lambda: _resolve_port( - "DYN_FORWARDPASS_METRIC_PORT", 20380 + "DYN_FORWARDPASS_METRIC_PORT", DEFAULT_FORWARDPASS_METRIC_PORT ), } diff --git a/components/src/dynamo/vllm/llm_engine.py b/components/src/dynamo/vllm/llm_engine.py index 69101366a2d9..93b35ed4c970 100644 --- a/components/src/dynamo/vllm/llm_engine.py +++ b/components/src/dynamo/vllm/llm_engine.py @@ -233,7 +233,7 @@ async def from_args( # don't re-parse (idempotent, but avoids a duplicate argparse + doubled # vLLM deprecation warnings at startup). if config is None: - config = parse_args(argv) + config = parse_args(argv, fpm_trace_relay_supported=False) if config.disaggregation_mode == DisaggregationMode.ENCODE: raise NotImplementedError( diff --git a/components/src/dynamo/vllm/main.py b/components/src/dynamo/vllm/main.py index bb5b5d257211..34b917f5c6e3 100644 --- a/components/src/dynamo/vllm/main.py +++ b/components/src/dynamo/vllm/main.py @@ -427,6 +427,7 @@ def setup_kv_event_publisher( def setup_fpm_relay( + config: Config, generate_endpoint: Endpoint, vllm_config: VllmConfig, ) -> Optional[list]: @@ -441,7 +442,7 @@ def setup_fpm_relay( Returns: List of FpmEventRelay instances, or None if FPM is not enabled. """ - if not envs.is_set("DYN_FORWARDPASS_METRIC_PORT"): + if not (envs.is_set("DYN_FORWARDPASS_METRIC_PORT") or config.fpm_trace): return None try: diff --git a/components/src/dynamo/vllm/tests/test_vllm_unified_main.py b/components/src/dynamo/vllm/tests/test_vllm_unified_main.py index 98c45c53ce4c..6739c09c8159 100644 --- a/components/src/dynamo/vllm/tests/test_vllm_unified_main.py +++ b/components/src/dynamo/vllm/tests/test_vllm_unified_main.py @@ -31,7 +31,7 @@ def test_main_routes_headless_to_run_dynamo_headless(): ): unified_main.main() - parse_args.assert_called_once() + parse_args.assert_called_once_with(fpm_trace_relay_supported=False) run_headless.assert_called_once_with(config) run.assert_not_called() @@ -47,7 +47,7 @@ async def test_main_routes_normal_node_to_run(): ): unified_main.main() - parse_args.assert_called_once() + parse_args.assert_called_once_with(fpm_trace_relay_supported=False) run_headless.assert_not_called() run.assert_called_once() assert run.call_args.args == (unified_main.VllmLLMEngine,) @@ -74,6 +74,8 @@ async def test_from_args_rejects_headless_reaching_engine_path(): headless=True, disaggregation_mode=DisaggregationMode.AGGREGATED, ) - with patch("dynamo.vllm.llm_engine.parse_args", return_value=config): + with patch("dynamo.vllm.llm_engine.parse_args", return_value=config) as parse_args: with pytest.raises(NotImplementedError, match="headless"): await VllmLLMEngine.from_args([]) + + parse_args.assert_called_once_with([], fpm_trace_relay_supported=False) diff --git a/components/src/dynamo/vllm/tests/test_vllm_unit.py b/components/src/dynamo/vllm/tests/test_vllm_unit.py index f1dd0021dd43..afe52d0b91fd 100644 --- a/components/src/dynamo/vllm/tests/test_vllm_unit.py +++ b/components/src/dynamo/vllm/tests/test_vllm_unit.py @@ -4,18 +4,23 @@ """Unit tests for vLLM backend components.""" import asyncio +import importlib import json +import logging +import os import re import socket import sys import warnings from contextlib import asynccontextmanager from pathlib import Path -from types import SimpleNamespace +from types import ModuleType, SimpleNamespace from unittest.mock import patch import pytest +import dynamo.llm as dynamo_llm +from dynamo.vllm import envs from dynamo.vllm.args import ( _connector_to_kv_transfer_json, _is_routable, @@ -57,6 +62,17 @@ mock_vllm_cli = make_cli_args_fixture("dynamo.vllm") +def _load_vllm_main() -> ModuleType: + """Load the entrypoint only in tests that need it. + + The lightweight pre-commit collection environment intentionally omits + uvloop, which ``dynamo.vllm.main`` imports at module scope. Eagerly loading + it here would prevent collection of every otherwise dependency-light test + in this module. + """ + return importlib.import_module("dynamo.vllm.main") + + def test_custom_jinja_template_invalid_path(mock_vllm_cli): """Test that invalid file path raises FileNotFoundError.""" invalid_path = "/nonexistent/path/to/template.jinja" @@ -350,8 +366,13 @@ def test_unified_from_args_applies_rl_logprobs_default(monkeypatch): dyn_reasoning_parser=None, ) worker_config = object() + parse_options = {} + + def fake_parse_args(argv, *, fpm_trace_relay_supported): + parse_options["fpm_trace_relay_supported"] = fpm_trace_relay_supported + return config - monkeypatch.setattr(llm_engine, "parse_args", lambda argv: config) + monkeypatch.setattr(llm_engine, "parse_args", fake_parse_args) monkeypatch.setattr( llm_engine.WorkerConfig, "from_runtime_config", @@ -366,6 +387,7 @@ async def run_from_args(): assert config.engine_args.logprobs_mode == "processed_logprobs" assert engine.enable_rl is True assert result_worker_config is worker_config + assert parse_options["fpm_trace_relay_supported"] is False def test_unified_generate_passes_enable_rl_to_sampling_params(monkeypatch): @@ -1077,6 +1099,11 @@ def _make_dynamo_config(**overrides): "disaggregation_mode": DisaggregationMode.AGGREGATED, "use_kv_events": False, "enable_local_indexer": True, + "embedding_worker": False, + "headless": False, + "multimodal_worker": False, + "multimodal_decode_worker": False, + "fpm_trace": False, "benchmark_mode": None, } defaults.update(overrides) @@ -1151,6 +1178,195 @@ def test_no_runner_attr_skipped_gracefully(self): assert not hasattr(engine_cfg, "runner") +class TestForwardPassMetricsActivation: + """FPM tracing should activate vLLM's existing FPM instrumentation.""" + + def test_cli_flag_enables_trace_and_exports_env(self, monkeypatch, mock_vllm_cli): + monkeypatch.delenv("DYN_FORWARDPASS_METRIC_PORT", raising=False) + monkeypatch.delenv("DYN_FPM_TRACE", raising=False) + mock_vllm_cli("--fpm-trace", "--model", "Qwen/Qwen3-0.6B") + + config = parse_args() + + assert config.fpm_trace is True + assert os.environ["DYN_FPM_TRACE"] == "1" + assert ( + config.engine_args.scheduler_cls + == "dynamo.vllm.instrumented_scheduler.InstrumentedScheduler" + ) + + def test_trace_enables_instrumented_scheduler_with_default_port(self, monkeypatch): + monkeypatch.delenv("DYN_FORWARDPASS_METRIC_PORT", raising=False) + dynamo_cfg = _make_dynamo_config(fpm_trace=True) + engine_cfg = _make_engine_config_with_runner(scheduler_cls=None) + + update_engine_config_with_dynamo(dynamo_cfg, engine_cfg) + + assert ( + engine_cfg.scheduler_cls + == "dynamo.vllm.instrumented_scheduler.InstrumentedScheduler" + ) + assert envs.DYN_FORWARDPASS_METRIC_PORT == 20380 + + def test_explicit_port_wins_when_trace_is_enabled(self, monkeypatch): + monkeypatch.setenv("DYN_FORWARDPASS_METRIC_PORT", "23456") + dynamo_cfg = _make_dynamo_config(fpm_trace=True) + engine_cfg = _make_engine_config_with_runner(scheduler_cls=None) + + update_engine_config_with_dynamo(dynamo_cfg, engine_cfg) + + assert ( + engine_cfg.scheduler_cls + == "dynamo.vllm.instrumented_scheduler.InstrumentedScheduler" + ) + assert envs.DYN_FORWARDPASS_METRIC_PORT == 23456 + + def test_false_trace_does_not_enable_fpm(self, monkeypatch): + monkeypatch.delenv("DYN_FORWARDPASS_METRIC_PORT", raising=False) + dynamo_cfg = _make_dynamo_config() + engine_cfg = _make_engine_config_with_runner(scheduler_cls=None) + + update_engine_config_with_dynamo(dynamo_cfg, engine_cfg) + + assert engine_cfg.scheduler_cls is None + + def test_disabled_trace_does_not_start_relay(self, monkeypatch): + vllm_main = _load_vllm_main() + monkeypatch.delenv("DYN_FORWARDPASS_METRIC_PORT", raising=False) + dynamo_cfg = _make_dynamo_config() + engine_cfg = _make_engine_config_with_runner(scheduler_cls=None) + + update_engine_config_with_dynamo(dynamo_cfg, engine_cfg) + assert ( + vllm_main.setup_fpm_relay(dynamo_cfg, SimpleNamespace(), SimpleNamespace()) + is None + ) + + assert engine_cfg.scheduler_cls is None + + def test_custom_scheduler_warns_and_serving_configuration_continues( + self, monkeypatch, caplog + ): + monkeypatch.delenv("DYN_FORWARDPASS_METRIC_PORT", raising=False) + dynamo_cfg = _make_dynamo_config(fpm_trace=True) + engine_cfg = _make_engine_config_with_runner( + scheduler_cls="example.CustomScheduler" + ) + + with caplog.at_level(logging.WARNING, logger="dynamo.vllm.args"): + update_engine_config_with_dynamo(dynamo_cfg, engine_cfg) + + assert engine_cfg.scheduler_cls == "example.CustomScheduler" + assert "InstrumentedScheduler will NOT be injected" in caplog.text + + def test_trace_only_starts_relay_on_default_port(self, monkeypatch): + vllm_main = _load_vllm_main() + monkeypatch.delenv("DYN_FORWARDPASS_METRIC_PORT", raising=False) + monkeypatch.setattr( + vllm_main, "get_dp_range_for_worker", lambda _config: (0, 1) + ) + + constructed = [] + + class FakeRelay: + def __init__(self, **kwargs): + constructed.append(kwargs) + + monkeypatch.setattr(dynamo_llm, "FpmEventRelay", FakeRelay, raising=False) + dynamo_cfg = _make_dynamo_config(fpm_trace=True) + endpoint = SimpleNamespace() + + relays = vllm_main.setup_fpm_relay(dynamo_cfg, endpoint, SimpleNamespace()) + + assert relays is not None + assert len(relays) == 1 + assert constructed == [ + { + "endpoint": endpoint, + "zmq_endpoint": "tcp://127.0.0.1:20380", + } + ] + + @pytest.mark.parametrize( + ("overrides", "role", "fpm_trace_relay_supported"), + [ + ({}, "unified backend", False), + ({"embedding_worker": True}, "embedding", True), + ({"headless": True}, "headless", True), + ( + {"disaggregation_mode": DisaggregationMode.ENCODE}, + "multimodal encode", + True, + ), + ], + ) + def test_trace_does_not_inject_scheduler_without_relay( + self, + monkeypatch, + caplog, + overrides, + role, + fpm_trace_relay_supported, + ): + monkeypatch.delenv("DYN_FORWARDPASS_METRIC_PORT", raising=False) + dynamo_cfg = _make_dynamo_config(fpm_trace=True, **overrides) + engine_cfg = _make_engine_config_with_runner(scheduler_cls=None) + + with caplog.at_level(logging.WARNING, logger="dynamo.vllm.args"): + update_engine_config_with_dynamo( + dynamo_cfg, + engine_cfg, + fpm_trace_relay_supported=fpm_trace_relay_supported, + ) + + assert engine_cfg.scheduler_cls is None + assert f"vLLM {role} workers do not create a Dynamo FPM relay" in caplog.text + + def test_explicit_port_preserves_legacy_activation_for_unsupported_role( + self, monkeypatch + ): + monkeypatch.setenv("DYN_FORWARDPASS_METRIC_PORT", "23456") + dynamo_cfg = _make_dynamo_config(embedding_worker=True) + engine_cfg = _make_engine_config_with_runner(scheduler_cls=None) + + update_engine_config_with_dynamo( + dynamo_cfg, + engine_cfg, + fpm_trace_relay_supported=False, + ) + + assert ( + engine_cfg.scheduler_cls + == "dynamo.vllm.instrumented_scheduler.InstrumentedScheduler" + ) + + def test_benchmark_does_not_reapply_trace_scheduler( + self, monkeypatch, caplog, tmp_path + ): + monkeypatch.delenv("DYN_FORWARDPASS_METRIC_PORT", raising=False) + dynamo_cfg = _make_dynamo_config( + fpm_trace=True, + benchmark_mode="agg", + benchmark_prefill_granularity=16, + benchmark_decode_length_granularity=6, + benchmark_decode_batch_granularity=6, + benchmark_warmup_iterations=5, + benchmark_output_path=str(tmp_path / "benchmark_results.json"), + benchmark_timeout=300, + ) + engine_cfg = _make_engine_config_with_runner(scheduler_cls=None) + + with caplog.at_level(logging.INFO, logger="dynamo.vllm.args"): + update_engine_config_with_dynamo(dynamo_cfg, engine_cfg) + + assert ( + engine_cfg.scheduler_cls + == "dynamo.vllm.instrumented_scheduler.InstrumentedScheduler" + ) + assert "Forward pass metrics enabled" in caplog.text + assert "Benchmark mode: auto-enabling InstrumentedScheduler" not in caplog.text + + class TestEmbeddingWorkerFlag: """Parsing + validation for --embedding-worker.""" diff --git a/components/src/dynamo/vllm/unified_main.py b/components/src/dynamo/vllm/unified_main.py index b2f4944417cf..f4ee13adf02d 100644 --- a/components/src/dynamo/vllm/unified_main.py +++ b/components/src/dynamo/vllm/unified_main.py @@ -22,7 +22,7 @@ def main(): # mp) run vLLM workers only and never touch the DistributedRuntime, so they # bypass the Worker/engine lifecycle entirely. Intercept before run() — # which would otherwise build the full backend and register an endpoint. - config = parse_args() + config = parse_args(fpm_trace_relay_supported=False) if config.headless: run_dynamo_headless(config) return diff --git a/components/src/dynamo/vllm/worker_factory.py b/components/src/dynamo/vllm/worker_factory.py index fe90dd46a494..7864d127bc45 100644 --- a/components/src/dynamo/vllm/worker_factory.py +++ b/components/src/dynamo/vllm/worker_factory.py @@ -513,7 +513,7 @@ async def _create_decode_worker( # Set up forward pass metrics relay (child ZMQ -> event plane). # In checkpoint mode the engine was created before the runtime, so # ForwardPassMetrics.worker_id will be empty (relay still works). - fpm_relays = self.setup_fpm_relay(generate_endpoint, vllm_config) + fpm_relays = self.setup_fpm_relay(config, generate_endpoint, vllm_config) if fpm_relays: handler.fpm_relays = fpm_relays @@ -755,7 +755,7 @@ async def _create_prefill_worker( # Set up forward pass metrics relay (child ZMQ -> event plane). # In checkpoint mode the engine was created before the runtime, so # ForwardPassMetrics.worker_id will be empty (relay still works). - fpm_relays = self.setup_fpm_relay(generate_endpoint, vllm_config) + fpm_relays = self.setup_fpm_relay(config, generate_endpoint, vllm_config) if fpm_relays: handler.fpm_relays = fpm_relays diff --git a/docs/index.yml b/docs/index.yml index 7a1c70b1a95b..5603e2f8ff8a 100644 --- a/docs/index.yml +++ b/docs/index.yml @@ -285,6 +285,8 @@ navigation: path: observability/health-checks.md - page: Tracing path: observability/tracing.md + - page: Forward Pass Metrics Tracing + path: observability/forward-pass-metrics-tracing.md - page: Logging path: observability/logging.md - section: Inference Simulation diff --git a/docs/observability/README.md b/docs/observability/README.md index 8e000e69189d..0e19a4aab8fc 100644 --- a/docs/observability/README.md +++ b/docs/observability/README.md @@ -59,6 +59,7 @@ For detailed setup instructions and configuration, see [Prometheus + Grafana Set | [Health Checks](health-checks.md) | Component health monitoring and readiness probes | `DYN_SYSTEM_PORT`†, `DYN_SYSTEM_STARTING_HEALTH_STATUS`, `DYN_SYSTEM_HEALTH_PATH`, `DYN_SYSTEM_LIVE_PATH`, `DYN_SYSTEM_USE_ENDPOINT_HEALTH_STATUS` | | [Tracing](tracing.md) | Distributed tracing with OpenTelemetry and Tempo | `DYN_LOGGING_JSONL`†, `OTEL_EXPORT_ENABLED`†, `OTEL_EXPORTER_OTLP_TRACES_ENDPOINT`†, `OTEL_SERVICE_NAME`† | | [Request Replay Tracing](request-tracing.md) | Per-request JSONL capture for direct DynoSim replay | `DYN_REQUEST_TRACE`, `DYN_REQUEST_TRACE_OUTPUT_PATH` | +| [Forward Pass Metrics Tracing](forward-pass-metrics-tracing.md) | Best-effort rotating gzip JSONL capture of backend forward pass metrics | `DYN_FPM_TRACE`, `DYN_FPM_OUTPUT_PATH`, `DYN_FPM_MODE`, `DYN_FPM_SAMPLE_INTERVAL_MS`, `DYN_FPM_JSONL_GZ_ROLL_BYTES`, `DYN_FPM_MAX_SEGMENTS` | | [Logging](logging.md) | Structured logging and OTLP log export to Loki | `DYN_LOGGING_JSONL`†, `DYN_LOG`, `DYN_LOG_USE_LOCAL_TZ`, `DYN_LOGGING_CONFIG_PATH`, `OTEL_SERVICE_NAME`†, `OTEL_EXPORT_ENABLED`†, `OTEL_EXPORTER_OTLP_TRACES_ENDPOINT`†, `OTEL_EXPORTER_OTLP_LOGS_ENDPOINT`† | | [Audit Payload Logging](logging.md#audit-payload-logging-otlp) | Per-request chat-completion payload capture exported over OTLP logs | `DYN_AUDIT_SINKS`, `DYN_AUDIT_FORCE_LOGGING`, `DYN_AUDIT_OTEL_MAX_PAYLOAD_BYTES`, `OTEL_EXPORTER_OTLP_LOGS_ENDPOINT`†, `OTEL_EXPORTER_OTLP_LOGS_PROTOCOL`† | diff --git a/docs/observability/forward-pass-metrics-tracing.md b/docs/observability/forward-pass-metrics-tracing.md new file mode 100644 index 000000000000..58c7eea23b40 --- /dev/null +++ b/docs/observability/forward-pass-metrics-tracing.md @@ -0,0 +1,233 @@ +--- +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +title: Forward Pass Metrics Tracing +subtitle: Persist backend forward pass metrics to rotating gzip JSONL files +--- + +Forward pass metrics (FPM) tracing is an opt-in, best-effort analysis stream. +It captures finalized FPM payloads in Dynamo's Rust publication path immediately +before they are sent to the event plane. This covers relay-backed vLLM and +SGLang publishers as well as the direct publishers used by TensorRT-LLM and the +mocker. Persistence is additive: it does not replace, suppress, or reroute those +events. On supported vLLM and SGLang worker topologies, however, +`--fpm-trace` (or its `DYN_FPM_TRACE=1` environment equivalent) also activates +the existing backend FPM generation and relay path, so enabling tracing can +cause FPM events to begin publishing. + +Enable tracing with the default sampled configuration: + +```bash +# CLI form +python -m dynamo.vllm --fpm-trace --model Qwen/Qwen3-0.6B + +# Equivalent environment form +export DYN_FPM_TRACE=1 +``` + +The shared runtime option also provides `--no-fpm-trace`, which overrides an +enabled `DYN_FPM_TRACE` value for that worker. + +Each worker writes its own files under `/tmp` by default: + +```text +/tmp/dynamo-fpm..000000.jsonl.gz +/tmp/dynamo-fpm..000001.jsonl.gz +``` + +`` is the worker's runtime connection ID, sanitized for use in a +file name. Partitioning files by producer prevents workers on a shared volume +from writing to the same gzip segment. + +> [!NOTE] +> `/tmp` is convenient but ephemeral. Mount a host directory or persistent +> volume and set `DYN_FPM_OUTPUT_PATH` if traces must survive pod replacement. + +## Configuration + +| Variable | Default when enabled | Description | +| --- | --- | --- | +| `DYN_FPM_TRACE` | unset | Environment form of the `--fpm-trace` / `--no-fpm-trace` switch. Accepts `1`/`0`, `true`/`false`, `on`/`off`, and `yes`/`no`, case-insensitively. | +| `DYN_FPM_OUTPUT_PATH` | `/tmp/dynamo-fpm` | Output prefix. Files are `...jsonl.gz`. | +| `DYN_FPM_MODE` | `sampled` | `sampled` keeps the latest changed record per worker and data-parallel rank; `full` captures every valid payload. | +| `DYN_FPM_SAMPLE_INTERVAL_MS` | `5000` | Positive sampling interval. Validated in both modes and used only in `sampled` mode. | +| `DYN_FPM_JSONL_GZ_ROLL_BYTES` | `268435456` | Positive uncompressed-byte threshold. Dynamo rolls before the next JSONL row would exceed it. | +| `DYN_FPM_MAX_SEGMENTS` | `4` | Positive number of segments retained per producer, including the active segment. | + +The other variables do not enable tracing by themselves. An invalid value or +an unwritable output path disables tracing only for that worker and produces a +warning. Inference and normal FPM publication continue. + +`--fpm-trace` or `DYN_FPM_TRACE=1` also enables the existing FPM generation +path for vLLM and SGLang. For vLLM, an explicit +`DYN_FORWARDPASS_METRIC_PORT` wins; otherwise Dynamo uses port `20380`. SGLang +uses its existing per-worker IPC endpoint. If vLLM has an incompatible custom +scheduler, Dynamo warns and continues serving without trace data from that +scheduler. + +TensorRT-LLM and the mocker already use Dynamo's direct FPM publisher on their +normal publication paths. For those publishers, `DYN_FPM_TRACE` adds local +persistence but does not need to activate a Python backend relay. + +Trace-based activation is limited to worker paths that construct a Dynamo FPM +relay: + +| Backend | Worker topology | `--fpm-trace` / `DYN_FPM_TRACE` support | +| --- | --- | --- | +| vLLM (`python -m dynamo.vllm`) | Aggregated, prefill, or decode, including native multimodal workers | Supported | +| vLLM (`python -m dynamo.vllm`) | Embedding, multimodal encode, or headless | Not supported; Dynamo warns and does not inject the FPM scheduler | +| vLLM (`python -m dynamo.vllm.unified_main`) | All worker topologies | Not supported; the unified path does not yet construct an FPM relay | +| SGLang (`python -m dynamo.sglang`) | Standard aggregated, prefill, decode, or LLM diffusion; this includes `--enable-multimodal` without a dedicated encoder | Supported | +| SGLang (`python -m dynamo.sglang`) | Embedding or the dedicated multimodal topology selected by `--dedicated-mm-encoder` (and its legacy worker flags) | Not supported; Dynamo warns and does not auto-enable FPM | +| SGLang (`python -m dynamo.sglang`) | Image diffusion or video generation | Not supported; these paths do not run the SGLang FPM publisher | +| SGLang (`python -m dynamo.sglang`) | Snapshot mode | Not supported; FPM is disabled during snapshot startup with a warning | +| SGLang (`python -m dynamo.sglang.unified_main`) | All worker topologies | Not supported; the unified path does not yet construct an FPM relay | + +`DYN_FORWARDPASS_METRIC_PORT` remains a separate legacy opt-in and takes +precedence when it is set, including on topologies where trace-based activation +is unsupported. It can enable backend FPM generation, but it does not add a +missing Dynamo relay, so an unsupported topology still does not produce a trace +file. + +`DYN_FPM_BENCHMARK_OUTPUT_PATH` remains a separate benchmark-only output and +is not used for live tracing. + +## Capture Modes + +In `sampled` mode, Dynamo retains only the newest pending payload for each +`(namespace, component, worker_id, dp_rank)` key. On each monotonic sampling +interval, it writes keys whose `counter_id` changed. It does not repeat an +unchanged counter. Pending values are flushed during graceful shutdown. + +In `full` mode, Dynamo writes every valid payload that reaches the producer, +including idle heartbeats. This mode can generate substantially more data. +It intentionally has no record-rate cap: the bounded queue limits memory and +protects inference, but it does not limit storage traffic. Each active writer +flushes a non-empty batch on a one-second interval and can flush sooner when +its 1 MiB buffer fills, so I/O on a shared volume scales with the number and +activity of producers. Use the default sampled mode or node-local storage +unless the shared filesystem has been sized for the expected full-mode load. + +Both modes are best effort. Producer enqueueing is nonblocking, and a bounded +in-process queue drops trace records instead of delaying inference. The worker +logs structured dropped-record counts when a trace consumer falls behind. + +## Record Shape + +Each line uses the shared gzip JSONL envelope. The `event` object is the FPM +trace record; `observed_at_unix_ms` is its absolute observation time. The +outer `timestamp` is milliseconds elapsed since that writer started. + +```json +{ + "timestamp": 1250, + "event": { + "schema": "dynamo.fpm.trace.v1", + "source": { + "namespace": "default", + "component": "backend", + "producer_id": "4192" + }, + "capture_mode": "sampled", + "observed_at_unix_ms": 1782777601250, + "fpm": { + "version": 1, + "worker_id": "4192", + "dp_rank": 0, + "counter_id": 42, + "wall_time": 0.025, + "scheduled_requests": { + "num_prefill_requests": 2, + "sum_prefill_tokens": 256, + "var_prefill_length": 100.0, + "sum_prefill_kv_tokens": 64, + "num_decode_requests": 3, + "sum_decode_kv_tokens": 1024, + "var_decode_kv_tokens": 50.0 + }, + "queued_requests": { + "num_prefill_requests": 1, + "sum_prefill_tokens": 128, + "var_prefill_length": 0.0, + "num_decode_requests": 0, + "sum_decode_kv_tokens": 0, + "var_decode_kv_tokens": 0.0 + } + } + } +} +``` + +The nested `fpm` object is the complete canonical payload. Use its +`worker_id`, `dp_rank`, and `counter_id` together when checking continuity. +A gap in `counter_id` can result from sampling, local queue pressure, an +upstream vLLM or SGLang ZMQ drop, a crash, or node loss. The trace is not a +durable event plane. + +## Read and Size Traces + +Decompress all segments for one output prefix in index order: + +```bash +gzip -cd /mnt/logs/dynamo-fpm.4192.*.jsonl.gz | jq -c '.event' +``` + +Replace `4192` with the producer ID in the file name. + +The roll threshold counts uncompressed JSONL bytes, while disk usage is the +compressed gzip size. Measure representative records for capacity planning: + +```bash +gzip -cd /mnt/logs/dynamo-fpm.4192.*.jsonl.gz | wc -c +gzip -cd /mnt/logs/dynamo-fpm.4192.*.jsonl.gz | wc -l +``` + +Estimate uncompressed bytes per day as: + +```text +average JSONL row bytes * records per second * 86400 +``` + +At the default five-second sampling interval, one continuously changing rank +writes about 17,280 periodic rows per day, plus a possible graceful-shutdown +flush. For example, 600-byte rows are about 10.4 MB per rank per day before +compression. In `full` mode, the rate follows the backend's forward passes: +the same 600-byte row at 10 forward passes per second is about 518 MB per rank +per day before compression. + +Rotation starts a new gzip file before adding a row that would cross the +configured threshold. A single oversized row is written intact to an otherwise +empty segment. After a new segment is written, Dynamo removes only the oldest +files that exactly match that producer's prefix. On restart, the next index is +one greater than the highest matching existing index, even if there are gaps. + +The segment limit applies independently to each producer. A shared persistent +volume can therefore contain up to `DYN_FPM_MAX_SEGMENTS` files for every +currently or previously used producer ID. Use an external lifecycle policy to +remove stale producer sets. + +## Kubernetes Storage + +Set the output prefix inside a mounted volume. Existing Dynamo environment and +volume-mount configuration is sufficient; no Dynamo Operator API change is +required. + +```yaml +env: + - name: DYN_FPM_TRACE + value: "1" + - name: DYN_FPM_OUTPUT_PATH + value: /var/log/dynamo/fpm +volumeMounts: + - name: fpm-traces + mountPath: /var/log/dynamo +volumes: + - name: fpm-traces + persistentVolumeClaim: + claimName: dynamo-fpm-traces +``` + +Graceful shutdown flushes records already accepted by the trace pipeline. +`SIGKILL`, node loss, full queues, and upstream transport loss can still lose +records. Treat the files as operational telemetry for analysis, not as input +for planner warm-start or replay. diff --git a/lib/llm/src/audit/sink.rs b/lib/llm/src/audit/sink.rs index 6b17c3dccdd2..37056415562c 100644 --- a/lib/llm/src/audit/sink.rs +++ b/lib/llm/src/audit/sink.rs @@ -153,6 +153,7 @@ impl JsonlGzipAuditSink { flush_interval: Duration::from_millis(policy.jsonl_flush_interval_ms.max(1)), roll_uncompressed_bytes: policy.jsonl_gz_roll_bytes, roll_lines: policy.jsonl_gz_roll_lines, + max_segments: None, }, ) .await @@ -309,6 +310,7 @@ mod tests { flush_interval: Duration::from_secs(60), roll_uncompressed_bytes: 1024 * 1024, roll_lines: None, + max_segments: None, }, ) .await diff --git a/lib/llm/src/fpm_publisher.rs b/lib/llm/src/fpm_publisher.rs index f33cc208bdf9..3734542b0758 100644 --- a/lib/llm/src/fpm_publisher.rs +++ b/lib/llm/src/fpm_publisher.rs @@ -32,6 +32,69 @@ const FPM_VERSION: i32 = 1; /// Matches Python `_FpmPublisherThread.HEARTBEAT_INTERVAL`. const IDLE_HEARTBEAT_INTERVAL: Duration = Duration::from_secs(1); +fn report_fpm_trace_init( + result: anyhow::Result>, +) -> Option { + match result { + Ok(trace) => trace, + Err(error) => { + tracing::warn!( + %error, + "FPM trace initialization failed; continuing without local persistence" + ); + None + } + } +} + +async fn init_fpm_trace(component: &Component) -> Option { + let namespace = component.namespace().name(); + let component_name = component.name().to_string(); + let producer_id = component.drt().connection_id().to_string(); + let runtime_id = component.drt().runtime().id().to_string(); + report_fpm_trace_init( + crate::fpm_trace::init_from_env_with_shutdown( + &runtime_id, + &namespace, + &component_name, + &producer_id, + component.drt().child_token(), + Some(component.drt().register_graceful_task()), + ) + .await, + ) +} + +fn tap_relay_fpm_with(payload: &bytes::Bytes, tap: F) +where + F: FnOnce(bytes::Bytes), +{ + tap(payload.clone()); +} + +fn tap_relay_fpm(payload: &bytes::Bytes, trace: Option<&crate::fpm_trace::FpmTrace>) { + if let Some(trace) = trace { + tap_relay_fpm_with(payload, |payload| { + trace.publish_payload(payload); + }); + } +} + +fn tap_direct_fpm_with(payload: &[u8], tap: F) +where + F: FnOnce(bytes::Bytes), +{ + tap(bytes::Bytes::copy_from_slice(payload)); +} + +fn tap_direct_fpm(payload: &[u8], trace: Option<&crate::fpm_trace::FpmTrace>) { + if let Some(trace) = trace { + tap_direct_fpm_with(payload, |payload| { + trace.publish_payload(payload); + }); + } +} + /// A relay that bridges ForwardPassMetrics from a local raw ZMQ PUB socket /// to the Dynamo event plane. pub struct FpmEventRelay { @@ -49,11 +112,13 @@ impl FpmEventRelay { let cancel = CancellationToken::new(); let cancel_clone = cancel.clone(); + let trace = rt.block_on(init_fpm_trace(&component)); + let publisher = rt.block_on(async { EventPublisher::for_component(&component, FPM_TOPIC).await })?; rt.spawn(async move { - Self::relay_loop(zmq_endpoint, publisher, cancel_clone).await; + Self::relay_loop(zmq_endpoint, publisher, cancel_clone, trace).await; }); Ok(Self { cancel }) @@ -68,6 +133,7 @@ impl FpmEventRelay { zmq_endpoint: String, publisher: EventPublisher, cancel: CancellationToken, + trace: Option, ) { let socket = match connect_sub_socket(&zmq_endpoint, None).await { Ok(socket) => socket, @@ -92,8 +158,9 @@ impl FpmEventRelay { let mut frames = multipart_message(frames); // ZMQ multipart: [topic, seq, payload] if frames.len() == 3 { - let payload = frames.swap_remove(2); - if let Err(e) = publisher.publish_bytes(payload).await { + let payload = bytes::Bytes::from(frames.swap_remove(2)); + tap_relay_fpm(&payload, trace.as_ref()); + if let Err(e) = publisher.publish_bytes_ref(&payload).await { tracing::warn!("FPM relay: event plane publish failed: {e}"); } } else { @@ -263,6 +330,7 @@ impl FpmDirectPublisher { let cancel = CancellationToken::new(); let publisher = EventPublisher::for_component(&component, FPM_TOPIC).await?; + let trace = init_fpm_trace(&component).await; // Shared channel: per-dp_rank tasks send snapshots here. A single publisher task // serializes them into a reusable buffer and preserves event-plane publish ordering. @@ -288,6 +356,7 @@ impl FpmDirectPublisher { pending.counter_id, ) { Ok(()) => { + tap_direct_fpm(&payload, trace.as_ref()); if let Err(e) = publisher.publish_bytes_ref(&payload).await { tracing::warn!("FPM direct publisher: event plane publish failed: {e}"); } @@ -387,8 +456,49 @@ impl Drop for FpmDirectPublisher { mod tests { use super::*; use serde::Deserialize; + use std::cell::{Cell, RefCell}; use std::collections::HashMap; + #[test] + fn fpm_trace_initialization_errors_are_soft() { + // Trace persistence is auxiliary. Reporting an initialization failure + // must not turn it into a constructor error for either FPM publisher. + assert!(report_fpm_trace_init(Err(anyhow::anyhow!("unwritable trace path"))).is_none()); + } + + #[test] + fn relay_and_direct_paths_tap_each_payload_exactly_once() { + let relay_calls = Cell::new(0); + let relay_payload = bytes::Bytes::from_static(b"relay"); + tap_relay_fpm_with(&relay_payload, |payload| { + assert_eq!(payload.as_ref(), b"relay"); + relay_calls.set(relay_calls.get() + 1); + }); + assert_eq!(relay_calls.get(), 1); + + let direct_calls = Cell::new(0); + tap_direct_fpm_with(b"direct", |payload| { + assert_eq!(payload.as_ref(), b"direct"); + direct_calls.set(direct_calls.get() + 1); + }); + assert_eq!(direct_calls.get(), 1); + } + + #[test] + fn trace_tap_precedes_and_survives_event_plane_failure() { + let steps = RefCell::new(Vec::new()); + let payload = bytes::Bytes::from_static(b"fpm"); + + tap_relay_fpm_with(&payload, |_| steps.borrow_mut().push("trace")); + let publish_result: anyhow::Result<()> = { + steps.borrow_mut().push("event-plane"); + Err(anyhow::anyhow!("publish failed")) + }; + + assert!(publish_result.is_err()); + assert_eq!(*steps.borrow(), ["trace", "event-plane"]); + } + /// Verify that serialize_fpm produces valid msgpack that round-trips /// through deserialization with the exact field names and values /// expected by the Python `ForwardPassMetrics` schema. diff --git a/lib/llm/src/fpm_trace/config.rs b/lib/llm/src/fpm_trace/config.rs new file mode 100644 index 000000000000..cda7548128ac --- /dev/null +++ b/lib/llm/src/fpm_trace/config.rs @@ -0,0 +1,283 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +use std::sync::OnceLock; + +use dynamo_runtime::config::environment_names::llm::fpm_trace as env_fpm_trace; + +pub const DEFAULT_OUTPUT_PATH: &str = "/tmp/dynamo-fpm"; +pub const DEFAULT_SAMPLE_INTERVAL_MS: u64 = 5_000; +pub const DEFAULT_JSONL_GZ_ROLL_BYTES: u64 = 256 * 1024 * 1024; +pub const DEFAULT_MAX_SEGMENTS: usize = 4; +pub(crate) const DEFAULT_CAPACITY: usize = 8_192; +pub(crate) const DEFAULT_JSONL_BUFFER_BYTES: usize = 1024 * 1024; +pub(crate) const DEFAULT_JSONL_FLUSH_INTERVAL_MS: u64 = 1_000; + +#[derive(Clone, Copy, Debug, Default, Eq, PartialEq, serde::Serialize)] +#[serde(rename_all = "lowercase")] +pub enum FpmTraceMode { + #[default] + Sampled, + Full, +} + +#[derive(Clone, Debug)] +pub struct FpmTracePolicy { + pub enabled: bool, + pub output_path: String, + pub mode: FpmTraceMode, + pub sample_interval_ms: u64, + pub jsonl_gz_roll_bytes: u64, + pub max_segments: usize, +} + +static POLICY: OnceLock = OnceLock::new(); + +impl Default for FpmTracePolicy { + fn default() -> Self { + Self { + enabled: false, + output_path: DEFAULT_OUTPUT_PATH.to_string(), + mode: FpmTraceMode::Sampled, + sample_interval_ms: DEFAULT_SAMPLE_INTERVAL_MS, + jsonl_gz_roll_bytes: DEFAULT_JSONL_GZ_ROLL_BYTES, + max_segments: DEFAULT_MAX_SEGMENTS, + } + } +} + +fn parse_bool(value: &str) -> anyhow::Result { + match value.trim().to_ascii_lowercase().as_str() { + "1" | "true" | "on" | "yes" => Ok(true), + "0" | "false" | "off" | "no" => Ok(false), + _ => anyhow::bail!( + "{} must be one of true/false, 1/0, on/off, or yes/no", + env_fpm_trace::DYN_FPM_TRACE + ), + } +} + +fn positive_integer_from_env(name: &str, default: T) -> anyhow::Result +where + T: Copy + From + PartialEq + std::str::FromStr, +{ + let Some(value) = std::env::var(name).ok() else { + return Ok(default); + }; + let parsed = value + .trim() + .parse::() + .map_err(|_| anyhow::anyhow!("{name} must be a positive integer"))?; + if parsed == T::from(0) { + anyhow::bail!("{name} must be greater than zero"); + } + Ok(parsed) +} + +fn load_enabled_policy() -> anyhow::Result { + let output_path = match std::env::var(env_fpm_trace::DYN_FPM_OUTPUT_PATH) { + Ok(value) if value.trim().is_empty() => { + anyhow::bail!("{} must not be empty", env_fpm_trace::DYN_FPM_OUTPUT_PATH) + } + Ok(value) => value.trim().to_string(), + Err(_) => DEFAULT_OUTPUT_PATH.to_string(), + }; + + let mode = match std::env::var(env_fpm_trace::DYN_FPM_MODE) { + Ok(value) if value.trim().eq_ignore_ascii_case("full") => FpmTraceMode::Full, + Ok(value) if value.trim().eq_ignore_ascii_case("sampled") => FpmTraceMode::Sampled, + Ok(_) => anyhow::bail!("{} must be sampled or full", env_fpm_trace::DYN_FPM_MODE), + Err(_) => FpmTraceMode::Sampled, + }; + + Ok(FpmTracePolicy { + enabled: true, + output_path, + mode, + sample_interval_ms: positive_integer_from_env( + env_fpm_trace::DYN_FPM_SAMPLE_INTERVAL_MS, + DEFAULT_SAMPLE_INTERVAL_MS, + )?, + jsonl_gz_roll_bytes: positive_integer_from_env( + env_fpm_trace::DYN_FPM_JSONL_GZ_ROLL_BYTES, + DEFAULT_JSONL_GZ_ROLL_BYTES, + )?, + max_segments: positive_integer_from_env( + env_fpm_trace::DYN_FPM_MAX_SEGMENTS, + DEFAULT_MAX_SEGMENTS, + )?, + }) +} + +fn load_from_env() -> FpmTracePolicy { + let enabled = match std::env::var(env_fpm_trace::DYN_FPM_TRACE) { + Err(_) => return FpmTracePolicy::default(), + Ok(value) => match parse_bool(&value) { + Ok(enabled) => enabled, + Err(error) => { + tracing::warn!(%error, "invalid FPM trace configuration; tracing disabled"); + return FpmTracePolicy::default(); + } + }, + }; + if !enabled { + return FpmTracePolicy::default(); + } + + load_enabled_policy().unwrap_or_else(|error| { + tracing::warn!(%error, "invalid FPM trace configuration; tracing disabled"); + FpmTracePolicy::default() + }) +} + +pub fn policy() -> &'static FpmTracePolicy { + POLICY.get_or_init(load_from_env) +} + +pub fn is_enabled() -> bool { + policy().enabled +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + #[serial_test::serial] + fn disabled_defaults_are_bounded_and_sampled() { + temp_env::with_vars( + [ + (env_fpm_trace::DYN_FPM_TRACE, None::<&str>), + (env_fpm_trace::DYN_FPM_OUTPUT_PATH, None), + (env_fpm_trace::DYN_FPM_MODE, None), + (env_fpm_trace::DYN_FPM_SAMPLE_INTERVAL_MS, None), + (env_fpm_trace::DYN_FPM_JSONL_GZ_ROLL_BYTES, None), + (env_fpm_trace::DYN_FPM_MAX_SEGMENTS, None), + ], + || { + let policy = load_from_env(); + assert!(!policy.enabled); + assert_eq!(policy.output_path, DEFAULT_OUTPUT_PATH); + assert_eq!(policy.mode, FpmTraceMode::Sampled); + assert_eq!(policy.sample_interval_ms, DEFAULT_SAMPLE_INTERVAL_MS); + assert_eq!(policy.jsonl_gz_roll_bytes, DEFAULT_JSONL_GZ_ROLL_BYTES); + assert_eq!(policy.max_segments, DEFAULT_MAX_SEGMENTS); + }, + ); + } + + #[test] + #[serial_test::serial] + fn accepts_full_mode_and_numeric_overrides() { + temp_env::with_vars( + [ + (env_fpm_trace::DYN_FPM_TRACE, Some("yes")), + (env_fpm_trace::DYN_FPM_OUTPUT_PATH, Some(" /var/log/fpm ")), + (env_fpm_trace::DYN_FPM_MODE, Some(" FULL ")), + (env_fpm_trace::DYN_FPM_SAMPLE_INTERVAL_MS, Some("250")), + (env_fpm_trace::DYN_FPM_JSONL_GZ_ROLL_BYTES, Some("4096")), + (env_fpm_trace::DYN_FPM_MAX_SEGMENTS, Some("7")), + ], + || { + let policy = load_from_env(); + assert!(policy.enabled); + assert_eq!(policy.output_path, "/var/log/fpm"); + assert_eq!(policy.mode, FpmTraceMode::Full); + assert_eq!(policy.sample_interval_ms, 250); + assert_eq!(policy.jsonl_gz_roll_bytes, 4096); + assert_eq!(policy.max_segments, 7); + }, + ); + } + + #[test] + #[serial_test::serial] + fn invalid_enabled_configuration_disables_trace_as_a_unit() { + temp_env::with_vars( + [ + (env_fpm_trace::DYN_FPM_TRACE, Some("true")), + (env_fpm_trace::DYN_FPM_MODE, Some("sometimes")), + (env_fpm_trace::DYN_FPM_SAMPLE_INTERVAL_MS, Some("0")), + (env_fpm_trace::DYN_FPM_JSONL_GZ_ROLL_BYTES, Some("bad")), + (env_fpm_trace::DYN_FPM_MAX_SEGMENTS, Some("0")), + ], + || { + let policy = load_from_env(); + assert!(!policy.enabled); + assert_eq!(policy.mode, FpmTraceMode::Sampled); + assert_eq!(policy.sample_interval_ms, DEFAULT_SAMPLE_INTERVAL_MS); + assert_eq!(policy.jsonl_gz_roll_bytes, DEFAULT_JSONL_GZ_ROLL_BYTES); + assert_eq!(policy.max_segments, DEFAULT_MAX_SEGMENTS); + }, + ); + } + + #[test] + #[serial_test::serial] + fn each_invalid_numeric_or_path_override_disables_trace() { + for (invalid_name, invalid_value) in [ + (env_fpm_trace::DYN_FPM_SAMPLE_INTERVAL_MS, "0"), + (env_fpm_trace::DYN_FPM_SAMPLE_INTERVAL_MS, "not-a-number"), + (env_fpm_trace::DYN_FPM_JSONL_GZ_ROLL_BYTES, "0"), + (env_fpm_trace::DYN_FPM_JSONL_GZ_ROLL_BYTES, "not-a-number"), + (env_fpm_trace::DYN_FPM_MAX_SEGMENTS, "0"), + (env_fpm_trace::DYN_FPM_MAX_SEGMENTS, "not-a-number"), + (env_fpm_trace::DYN_FPM_OUTPUT_PATH, " "), + ] { + let mut vars = vec![ + (env_fpm_trace::DYN_FPM_TRACE, Some("true")), + (env_fpm_trace::DYN_FPM_OUTPUT_PATH, None), + (env_fpm_trace::DYN_FPM_MODE, None), + (env_fpm_trace::DYN_FPM_SAMPLE_INTERVAL_MS, None), + (env_fpm_trace::DYN_FPM_JSONL_GZ_ROLL_BYTES, None), + (env_fpm_trace::DYN_FPM_MAX_SEGMENTS, None), + ]; + let (_, value) = vars + .iter_mut() + .find(|(name, _)| *name == invalid_name) + .unwrap(); + *value = Some(invalid_value); + + temp_env::with_vars(vars, || { + assert!( + !load_from_env().enabled, + "invalid setting should disable trace: {invalid_name}={invalid_value}" + ); + }); + } + } + + #[test] + #[serial_test::serial] + fn invalid_master_switch_disables_trace() { + temp_env::with_var(env_fpm_trace::DYN_FPM_TRACE, Some("maybe"), || { + assert!(!load_from_env().enabled); + }); + } + + #[test] + #[serial_test::serial] + fn explicit_false_ignores_other_trace_settings() { + temp_env::with_vars( + [ + (env_fpm_trace::DYN_FPM_TRACE, Some("off")), + (env_fpm_trace::DYN_FPM_MODE, Some("invalid")), + (env_fpm_trace::DYN_FPM_MAX_SEGMENTS, Some("0")), + ], + || assert!(!load_from_env().enabled), + ); + } + + #[test] + fn master_switch_parser_accepts_only_documented_boolean_forms() { + for value in ["true", "TRUE", "1", "on", "ON", "yes"] { + assert!(parse_bool(value).unwrap(), "value={value}"); + } + for value in ["false", "FALSE", "0", "off", "OFF", "no"] { + assert!(!parse_bool(value).unwrap(), "value={value}"); + } + for value in ["", "enabled", "2"] { + assert!(parse_bool(value).is_err(), "value={value}"); + } + } +} diff --git a/lib/llm/src/fpm_trace/mod.rs b/lib/llm/src/fpm_trace/mod.rs new file mode 100644 index 000000000000..5c10896002e6 --- /dev/null +++ b/lib/llm/src/fpm_trace/mod.rs @@ -0,0 +1,572 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +//! Best-effort producer-side persistence for forward-pass metrics. +//! +//! This trace is an analysis aid, not a durable replacement for the FPM event +//! plane. Each producer publishes into its own bounded in-process queue and +//! never waits for disk I/O. A slow sink can therefore drop trace records +//! without delaying inference or normal FPM delivery. + +pub mod config; +mod sink; + +use std::collections::HashMap; +use std::sync::atomic::{AtomicBool, AtomicU64, AtomicUsize, Ordering}; +use std::sync::{Arc, OnceLock, Weak}; +use std::time::{SystemTime, UNIX_EPOCH}; + +use bytes::Bytes; +use tokio::sync::{Mutex, Notify, broadcast}; +use tokio_util::sync::CancellationToken; + +use dynamo_runtime::utils::GracefulTaskGuard; + +pub use config::{FpmTraceMode, FpmTracePolicy, is_enabled, policy}; + +#[derive(Clone, Debug)] +pub(crate) struct FpmTraceEvent { + observed_at_unix_ms: u64, + payload: Bytes, + source: Arc, +} + +#[derive(Clone, Debug, Eq, Hash, PartialEq)] +struct FpmTraceKey { + runtime_id: String, + producer_id: String, +} + +struct FpmTraceInner { + producer_id: String, + sender: broadcast::Sender, + shutdown: CancellationToken, + accepting: AtomicBool, + publishers_in_flight: AtomicUsize, + publishers_idle: Notify, + leases: AtomicUsize, + dropped: AtomicU64, + closed: AtomicBool, + closed_notify: Notify, +} + +impl FpmTraceInner { + fn begin_publish(&self) -> bool { + if !self.accepting.load(Ordering::SeqCst) { + return false; + } + self.publishers_in_flight.fetch_add(1, Ordering::SeqCst); + if self.accepting.load(Ordering::SeqCst) { + true + } else { + self.finish_publish(); + false + } + } + + fn finish_publish(&self) { + if self.publishers_in_flight.fetch_sub(1, Ordering::SeqCst) == 1 { + self.publishers_idle.notify_waiters(); + } + } + + async fn stop_accepting(&self) { + self.accepting.store(false, Ordering::SeqCst); + loop { + let idle = self.publishers_idle.notified(); + if self.publishers_in_flight.load(Ordering::SeqCst) == 0 { + return; + } + idle.await; + } + } + + fn record_dropped(&self, count: u64) -> u64 { + self.dropped.fetch_add(count, Ordering::Relaxed) + count + } + + fn mark_closed(&self) { + self.accepting.store(false, Ordering::SeqCst); + self.closed.store(true, Ordering::Release); + self.publishers_idle.notify_waiters(); + self.closed_notify.notify_waiters(); + } + + async fn wait_closed(&self) { + loop { + let notified = self.closed_notify.notified(); + if self.closed.load(Ordering::Acquire) { + return; + } + notified.await; + } + } +} + +/// Producer-owned, nonblocking handle to one FPM trace sink. +/// +/// Handles with the same runtime and producer share one queue and one writer. +/// Each handle stamps its own namespace/component source on every event, so +/// components sharing a producer cannot collide on the output file or inherit +/// another component's attribution. Different producers have independent +/// queues, files, and retention budgets. +pub(crate) struct FpmTrace { + inner: Arc, + source: Arc, +} + +impl FpmTrace { + fn try_new_lease(inner: Arc, source: Arc) -> Option { + // Increment only while at least one producer handle is still alive. + // This makes lease acquisition atomic with the last handle's 1 -> 0 + // transition, so a concurrent initializer cannot resurrect a sink + // after its shutdown has begun. + inner + .leases + .fetch_update(Ordering::AcqRel, Ordering::Acquire, |leases| { + if leases == 0 { + None + } else { + leases.checked_add(1) + } + }) + .ok()?; + Some(Self { inner, source }) + } + + /// Try to enqueue one finalized FPM msgpack payload for persistence. + /// + /// This never waits. A bounded broadcast queue intentionally retains the + /// newest records under load; the sink reports overwritten-record counts. + /// `false` means shutdown had already stopped accepting trace records. + pub(crate) fn publish_payload(&self, payload: Bytes) -> bool { + if !self.inner.begin_publish() { + return false; + } + let observed_at_unix_ms = SystemTime::now() + .duration_since(UNIX_EPOCH) + .map(|duration| duration.as_millis() as u64) + .unwrap_or_default(); + let sent = self + .inner + .sender + .send(FpmTraceEvent { + observed_at_unix_ms, + payload, + source: self.source.clone(), + }) + .is_ok(); + self.inner.finish_publish(); + if !sent { + self.inner.record_dropped(1); + } + sent + } + + #[cfg(test)] + async fn wait_closed(&self) { + self.inner.wait_closed().await; + } +} + +impl Clone for FpmTrace { + fn clone(&self) -> Self { + Self::try_new_lease(self.inner.clone(), self.source.clone()) + .expect("a live FPM trace handle must own a producer lease") + } +} + +impl Drop for FpmTrace { + fn drop(&mut self) { + if self.inner.leases.fetch_sub(1, Ordering::AcqRel) == 1 { + // The last producer stopped. Close its sink even when the runtime + // itself remains alive; other producers have independent sinks. + self.inner.shutdown.cancel(); + } + } +} + +enum TraceRegistryEntry { + Empty, + Active(Weak), + /// Initialization failures are terminal for this producer. A multi-DP + /// producer otherwise retries the same bad path and emits one warning per + /// relay, while failures for unrelated producers must remain isolated. + Failed, +} + +type TraceSlot = Arc>; +type TraceRegistry = HashMap; +static REGISTRY: OnceLock> = OnceLock::new(); + +fn registry() -> &'static Mutex { + REGISTRY.get_or_init(|| Mutex::new(HashMap::new())) +} + +pub(crate) async fn init_from_env_with_shutdown( + runtime_id: &str, + namespace: &str, + component: &str, + producer_id: &str, + shutdown: CancellationToken, + graceful_guard: Option, +) -> anyhow::Result> { + let policy = policy(); + if !policy.enabled { + return Ok(None); + } + + init_with_policy( + runtime_id, + namespace, + component, + producer_id, + policy.clone(), + shutdown, + graceful_guard, + ) + .await +} + +async fn init_with_policy( + runtime_id: &str, + namespace: &str, + component: &str, + producer_id: &str, + policy: FpmTracePolicy, + shutdown: CancellationToken, + graceful_guard: Option, +) -> anyhow::Result> { + let source = sink::FpmTraceSource { + namespace: namespace.to_string(), + component: component.to_string(), + producer_id: producer_id.to_string(), + }; + let source = Arc::new(source); + let key = FpmTraceKey { + runtime_id: runtime_id.to_string(), + producer_id: producer_id.to_string(), + }; + + // Serialize initialization only within one producer. Unrelated producers + // can preflight and open their writers concurrently. + let slot = { + let mut registry = registry().lock().await; + registry + .entry(key) + .or_insert_with(|| Arc::new(Mutex::new(TraceRegistryEntry::Empty))) + .clone() + }; + let mut entry = slot.lock().await; + match &*entry { + TraceRegistryEntry::Failed => return Ok(None), + TraceRegistryEntry::Active(inner) => { + if let Some(inner) = inner.upgrade() { + if let Some(trace) = FpmTrace::try_new_lease(inner.clone(), source.clone()) { + if !inner.shutdown.is_cancelled() && !inner.closed.load(Ordering::Acquire) { + return Ok(Some(trace)); + } + drop(trace); + } + // Do not let a rapid producer restart open the same segment + // while the previous gzip writer is still finalizing it. + inner.wait_closed().await; + } + } + TraceRegistryEntry::Empty => {} + } + + let (sender, receiver) = broadcast::channel(config::DEFAULT_CAPACITY); + let inner = Arc::new(FpmTraceInner { + producer_id: producer_id.to_string(), + sender, + // A child token makes producer-local close idempotent while also + // inheriting runtime phase-1 cancellation. + shutdown: shutdown.child_token(), + accepting: AtomicBool::new(true), + publishers_in_flight: AtomicUsize::new(0), + publishers_idle: Notify::new(), + // The handle returned below is the initial producer lease. + leases: AtomicUsize::new(1), + dropped: AtomicU64::new(0), + closed: AtomicBool::new(false), + closed_notify: Notify::new(), + }); + + if let Err(error) = + sink::spawn_worker(policy.clone(), receiver, inner.clone(), graceful_guard).await + { + *entry = TraceRegistryEntry::Failed; + return Err(error); + } + *entry = TraceRegistryEntry::Active(Arc::downgrade(&inner)); + + tracing::info!( + namespace, + component, + producer_id, + mode = ?policy.mode, + sample_interval_ms = policy.sample_interval_ms, + output_path = %policy.output_path, + max_segments = policy.max_segments, + "FPM trace initialized" + ); + Ok(Some(FpmTrace { inner, source })) +} + +#[cfg(test)] +mod tests { + use std::io::Read; + use std::path::Path; + + use flate2::read::MultiGzDecoder; + use serde_json::Value; + + use super::*; + use crate::telemetry::jsonl_gz::segment_path; + + fn full_policy(path: &Path) -> FpmTracePolicy { + FpmTracePolicy { + enabled: true, + output_path: path.display().to_string(), + mode: FpmTraceMode::Full, + sample_interval_ms: 100, + jsonl_gz_roll_bytes: 1024 * 1024, + max_segments: 4, + } + } + + fn payload(worker_id: &str, counter_id: i64) -> Bytes { + Bytes::from( + rmp_serde::to_vec_named(&serde_json::json!({ + "version": 1, + "worker_id": worker_id, + "dp_rank": 0, + "counter_id": counter_id, + "wall_time": 0.01, + "scheduled_requests": { + "num_prefill_requests": 1, + "sum_prefill_tokens": 16, + "var_prefill_length": 0.0, + "sum_prefill_kv_tokens": 0, + "num_decode_requests": 2, + "sum_decode_kv_tokens": 64, + "var_decode_kv_tokens": 0.0, + }, + "queued_requests": { + "num_prefill_requests": 0, + "sum_prefill_tokens": 0, + "var_prefill_length": 0.0, + "num_decode_requests": 0, + "sum_decode_kv_tokens": 0, + "var_decode_kv_tokens": 0.0, + }, + })) + .unwrap(), + ) + } + + fn records(base_path: &Path, producer_id: &str) -> Vec { + let producer_path = + sink::producer_output_path(&base_path.display().to_string(), producer_id); + let bytes = std::fs::read(segment_path(Path::new(&producer_path), 0)).unwrap(); + let mut decoder = MultiGzDecoder::new(bytes.as_slice()); + let mut content = String::new(); + decoder.read_to_string(&mut content).unwrap(); + content + .lines() + .map(|line| serde_json::from_str(line).unwrap()) + .collect() + } + + fn trace_without_worker() -> FpmTrace { + let (sender, _receiver) = broadcast::channel(1); + let inner = Arc::new(FpmTraceInner { + producer_id: "producer-lease-test".to_string(), + sender, + shutdown: CancellationToken::new(), + accepting: AtomicBool::new(true), + publishers_in_flight: AtomicUsize::new(0), + publishers_idle: Notify::new(), + leases: AtomicUsize::new(1), + dropped: AtomicU64::new(0), + closed: AtomicBool::new(false), + closed_notify: Notify::new(), + }); + let source = Arc::new(sink::FpmTraceSource { + namespace: "namespace".to_string(), + component: "backend".to_string(), + producer_id: "producer-lease-test".to_string(), + }); + FpmTrace { inner, source } + } + + #[test] + fn producer_lease_cannot_be_resurrected_after_last_handle_drops() { + let trace = trace_without_worker(); + let inner = trace.inner.clone(); + let source = trace.source.clone(); + + drop(trace); + + assert!(inner.shutdown.is_cancelled()); + assert_eq!(inner.leases.load(Ordering::Acquire), 0); + assert!(FpmTrace::try_new_lease(inner.clone(), source).is_none()); + assert_eq!(inner.leases.load(Ordering::Acquire), 0); + } + + #[tokio::test] + async fn producer_registry_shares_matching_producers_and_isolates_others() { + let dir = tempfile::tempdir().unwrap(); + let base_path = dir.path().join("trace"); + let policy = full_policy(&base_path); + let runtime_a_shutdown = CancellationToken::new(); + let runtime_b_shutdown = CancellationToken::new(); + + let producer_a = init_with_policy( + "runtime-a", + "namespace", + "backend", + "producer-a", + policy.clone(), + runtime_a_shutdown.clone(), + None, + ) + .await + .unwrap() + .unwrap(); + let producer_a_second_relay = init_with_policy( + "runtime-a", + "namespace", + "backend", + "producer-a", + policy.clone(), + runtime_a_shutdown.clone(), + None, + ) + .await + .unwrap() + .unwrap(); + let producer_a_other_component = init_with_policy( + "runtime-a", + "namespace", + "sidecar", + "producer-a", + policy.clone(), + runtime_a_shutdown.clone(), + None, + ) + .await + .unwrap() + .unwrap(); + let producer_b = init_with_policy( + "runtime-b", + "namespace", + "backend", + "producer-b", + policy, + runtime_b_shutdown.clone(), + None, + ) + .await + .unwrap() + .unwrap(); + + assert!(Arc::ptr_eq( + &producer_a.inner, + &producer_a_second_relay.inner + )); + assert!(Arc::ptr_eq( + &producer_a.inner, + &producer_a_other_component.inner + )); + assert!(!Arc::ptr_eq(&producer_a.inner, &producer_b.inner)); + + assert!(producer_a.publish_payload(payload("worker-a", 1))); + assert!(producer_a_second_relay.publish_payload(payload("worker-a", 2))); + assert!(producer_a_other_component.publish_payload(payload("worker-a", 3))); + assert!(producer_b.publish_payload(payload("worker-b", 7))); + + // Runtime shutdown is idempotent. Each worker first closes the + // publish acceptance boundary, then drains the stable queue. + runtime_a_shutdown.cancel(); + runtime_a_shutdown.cancel(); + runtime_b_shutdown.cancel(); + runtime_b_shutdown.cancel(); + producer_a.wait_closed().await; + producer_b.wait_closed().await; + + let a_records = records(&base_path, "producer-a"); + let b_records = records(&base_path, "producer-b"); + let a_counters: Vec<_> = a_records + .iter() + .map(|record| record["event"]["fpm"]["counter_id"].as_i64().unwrap()) + .collect(); + assert_eq!(a_counters, [1, 2, 3]); + assert_eq!(b_records.len(), 1); + assert_eq!(b_records[0]["event"]["fpm"]["counter_id"], 7); + assert!(a_records[..2].iter().all(|record| { + record["event"]["source"]["producer_id"] == "producer-a" + && record["event"]["fpm"]["worker_id"] == "worker-a" + && record["event"]["source"]["component"] == "backend" + })); + assert_eq!(a_records[2]["event"]["source"]["component"], "sidecar"); + assert_eq!(b_records[0]["event"]["source"]["producer_id"], "producer-b"); + assert!(!producer_a.publish_payload(payload("worker-a", 3))); + } + + #[tokio::test] + async fn failed_initialization_is_terminal_only_for_its_producer() { + let dir = tempfile::tempdir().unwrap(); + let blocked_parent = dir.path().join("blocked-parent"); + std::fs::write(&blocked_parent, b"not a directory").unwrap(); + let failed_policy = full_policy(&blocked_parent.join("trace")); + let valid_base = dir.path().join("valid-trace"); + let valid_policy = full_policy(&valid_base); + + let first = init_with_policy( + "failure-runtime", + "namespace", + "backend", + "failed-producer", + failed_policy, + CancellationToken::new(), + None, + ) + .await; + assert!(first.is_err()); + + // A valid policy would succeed if retried. The cached failure instead + // returns disabled without repeating preflight or another warning. + let repeated = init_with_policy( + "failure-runtime", + "namespace", + "backend", + "failed-producer", + valid_policy.clone(), + CancellationToken::new(), + None, + ) + .await + .unwrap(); + assert!(repeated.is_none()); + + let other_shutdown = CancellationToken::new(); + let other = init_with_policy( + "failure-runtime", + "namespace", + "backend", + "healthy-producer", + valid_policy, + other_shutdown.clone(), + None, + ) + .await + .unwrap() + .expect("one producer's failure must not disable another producer"); + assert!(other.publish_payload(payload("healthy-worker", 1))); + other_shutdown.cancel(); + other.wait_closed().await; + assert_eq!(records(&valid_base, "healthy-producer").len(), 1); + } +} diff --git a/lib/llm/src/fpm_trace/sink.rs b/lib/llm/src/fpm_trace/sink.rs new file mode 100644 index 000000000000..3f4e1df1c5f5 --- /dev/null +++ b/lib/llm/src/fpm_trace/sink.rs @@ -0,0 +1,846 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +use std::collections::BTreeMap; +use std::io::Write as _; +use std::path::Path; +use std::sync::Arc; +use std::sync::atomic::{AtomicU64, Ordering}; +use std::time::Duration; + +use anyhow::Context as _; +use serde::Serialize; +use serde_json::{Map, Value}; +use tokio::sync::broadcast; + +use dynamo_runtime::utils::GracefulTaskGuard; + +use crate::telemetry::jsonl_gz::{JsonlGzipSinkOptions, JsonlGzipWriter}; + +use super::config::{ + DEFAULT_JSONL_BUFFER_BYTES, DEFAULT_JSONL_FLUSH_INTERVAL_MS, FpmTraceMode, FpmTracePolicy, +}; +use super::{FpmTraceEvent, FpmTraceInner}; + +static PROBE_SEQUENCE: AtomicU64 = AtomicU64::new(0); + +#[derive(Clone, Debug, Eq, Hash, Ord, PartialEq, PartialOrd, Serialize)] +pub(super) struct FpmTraceSource { + pub(super) namespace: String, + pub(super) component: String, + pub(super) producer_id: String, +} + +#[derive(Clone, Debug, Serialize)] +struct FpmTraceRecord { + schema: &'static str, + source: FpmTraceSource, + capture_mode: FpmTraceMode, + observed_at_unix_ms: u64, + fpm: Value, +} + +// A producer can expose multiple Dynamo components through one shared trace +// worker. Keep their sampled state independent even when backend worker IDs +// and data-parallel ranks overlap. +type FpmKey = (FpmTraceSource, String, i64); + +fn require_nonnegative_integer(object: &Map, field: &str) -> anyhow::Result<()> { + let valid = object.get(field).is_some_and(|value| { + value.as_u64().is_some() || value.as_i64().is_some_and(|number| number >= 0) + }); + if !valid { + anyhow::bail!("FPM field {field} must be a non-negative integer"); + } + Ok(()) +} + +fn require_nonnegative_finite_number( + object: &Map, + field: &str, +) -> anyhow::Result<()> { + let valid = object + .get(field) + .and_then(Value::as_f64) + .is_some_and(|number| number.is_finite() && number >= 0.0); + if !valid { + anyhow::bail!("FPM field {field} must be a finite non-negative number"); + } + Ok(()) +} + +fn validate_request_metrics( + object: &Map, + integer_fields: &[&str], + variance_fields: &[&str], +) -> anyhow::Result<()> { + for field in integer_fields { + require_nonnegative_integer(object, field)?; + } + for field in variance_fields { + require_nonnegative_finite_number(object, field)?; + } + Ok(()) +} + +fn validate_canonical_fpm(object: &Map) -> anyhow::Result<()> { + if object.get("version").and_then(Value::as_i64) != Some(1) { + anyhow::bail!("FPM payload has unsupported or missing version"); + } + require_nonnegative_finite_number(object, "wall_time")?; + + let scheduled = object + .get("scheduled_requests") + .and_then(Value::as_object) + .ok_or_else(|| anyhow::anyhow!("FPM payload has no scheduled_requests map"))?; + validate_request_metrics( + scheduled, + &[ + "num_prefill_requests", + "sum_prefill_tokens", + "sum_prefill_kv_tokens", + "num_decode_requests", + "sum_decode_kv_tokens", + ], + &["var_prefill_length", "var_decode_kv_tokens"], + )?; + + let queued = object + .get("queued_requests") + .and_then(Value::as_object) + .ok_or_else(|| anyhow::anyhow!("FPM payload has no queued_requests map"))?; + validate_request_metrics( + queued, + &[ + "num_prefill_requests", + "sum_prefill_tokens", + "num_decode_requests", + "sum_decode_kv_tokens", + ], + &["var_prefill_length", "var_decode_kv_tokens"], + )?; + Ok(()) +} + +fn decode_event( + event: FpmTraceEvent, + capture_mode: FpmTraceMode, +) -> anyhow::Result<(FpmKey, i64, FpmTraceRecord)> { + let fpm: Value = rmp_serde::from_slice(&event.payload).context("decoding FPM msgpack")?; + let object = fpm + .as_object() + .ok_or_else(|| anyhow::anyhow!("FPM payload is not a msgpack map"))?; + validate_canonical_fpm(object)?; + let worker_id = object + .get("worker_id") + .and_then(Value::as_str) + .ok_or_else(|| anyhow::anyhow!("FPM payload has no string worker_id"))? + .to_string(); + let dp_rank = object + .get("dp_rank") + .and_then(Value::as_i64) + .filter(|dp_rank| *dp_rank >= 0) + .ok_or_else(|| anyhow::anyhow!("FPM payload has no non-negative integer dp_rank"))?; + let counter_id = object + .get("counter_id") + .and_then(Value::as_i64) + .filter(|counter_id| *counter_id >= 0) + .ok_or_else(|| anyhow::anyhow!("FPM payload has no non-negative integer counter_id"))?; + + let source = (*event.source).clone(); + Ok(( + (source.clone(), worker_id, dp_rank), + counter_id, + FpmTraceRecord { + schema: "dynamo.fpm.trace.v1", + source, + capture_mode, + observed_at_unix_ms: event.observed_at_unix_ms, + fpm, + }, + )) +} + +fn sanitize_producer_id(producer_id: &str) -> String { + let sanitized: String = producer_id + .chars() + .map(|character| { + if character.is_ascii_alphanumeric() || matches!(character, '-' | '_') { + character + } else { + '_' + } + }) + .collect(); + if sanitized.is_empty() { + "unknown".to_string() + } else { + sanitized + } +} + +pub(super) fn producer_output_path(base_path: &str, producer_id: &str) -> String { + let prefix = base_path + .strip_suffix(".jsonl.gz") + .or_else(|| base_path.strip_suffix(".jsonl")) + .unwrap_or(base_path); + format!("{prefix}.{}", sanitize_producer_id(producer_id)) +} + +fn preflight_writable_parent(output_path: &str) -> anyhow::Result<()> { + let output_path = Path::new(output_path); + let parent = output_path + .parent() + .filter(|parent| !parent.as_os_str().is_empty()) + .unwrap_or_else(|| Path::new(".")); + std::fs::create_dir_all(parent) + .with_context(|| format!("creating FPM trace directory {}", parent.display()))?; + + let sequence = PROBE_SEQUENCE.fetch_add(1, Ordering::Relaxed); + let producer_prefix = output_path + .file_name() + .and_then(|name| name.to_str()) + .unwrap_or("dynamo-fpm"); + let probe_path = parent.join(format!( + ".{producer_prefix}.write-probe-{}-{sequence}", + std::process::id() + )); + let mut probe = std::fs::OpenOptions::new() + .create_new(true) + .write(true) + .open(&probe_path) + .with_context(|| format!("creating FPM trace write probe {}", probe_path.display()))?; + + let write_result = probe + .write_all(b"dynamo-fpm-write-probe") + .and_then(|_| probe.sync_all()); + drop(probe); + if let Err(error) = write_result { + let _ = std::fs::remove_file(&probe_path); + return Err(error) + .with_context(|| format!("writing FPM trace probe {}", probe_path.display())); + } + std::fs::remove_file(&probe_path) + .with_context(|| format!("removing FPM trace write probe {}", probe_path.display()))?; + Ok(()) +} + +pub(super) async fn spawn_worker( + policy: FpmTracePolicy, + receiver: broadcast::Receiver, + owner: Arc, + graceful_guard: Option, +) -> anyhow::Result<()> { + let output_path = producer_output_path(&policy.output_path, &owner.producer_id); + let preflight_path = output_path.clone(); + tokio::task::spawn_blocking(move || preflight_writable_parent(&preflight_path)) + .await + .context("joining FPM trace write preflight")??; + let writer = JsonlGzipWriter::new( + output_path.clone(), + JsonlGzipSinkOptions { + buffer_bytes: DEFAULT_JSONL_BUFFER_BYTES, + flush_interval: Duration::from_millis(DEFAULT_JSONL_FLUSH_INTERVAL_MS), + roll_uncompressed_bytes: policy.jsonl_gz_roll_bytes, + roll_lines: None, + max_segments: Some(policy.max_segments), + }, + ) + .await + .with_context(|| format!("opening FPM gzip JSONL trace at {output_path}"))?; + let completion = MarkClosedOnDrop(owner.clone()); + tokio::spawn(async move { + // The guard keeps runtime shutdown in phase 2 until this task has + // drained the queue and awaited writer close. + let _graceful_guard = graceful_guard; + // Constructed outside the async block, so task abort before its first + // poll still marks the producer closed and wakes registry waiters. It + // is declared after the graceful guard so it drops first on return. + let _completion = completion; + // Keeping the owner alive keeps this producer discoverable in the + // registry until its writer has finished closing. + match policy.mode { + FpmTraceMode::Full => run_full(receiver, writer, owner.clone()).await, + FpmTraceMode::Sampled => { + run_sampled(receiver, writer, owner.clone(), policy.sample_interval_ms).await + } + } + }); + + Ok(()) +} + +struct MarkClosedOnDrop(Arc); + +impl Drop for MarkClosedOnDrop { + fn drop(&mut self) { + self.0.mark_closed(); + } +} + +async fn send_decoded(writer: &JsonlGzipWriter, event: FpmTraceEvent) { + match decode_event(event, FpmTraceMode::Full) { + Ok((_, _, record)) => { + if writer.send(record).await.is_err() { + tracing::warn!("FPM trace writer closed; dropping record"); + } + } + Err(error) => tracing::warn!(%error, "FPM trace dropped malformed payload"), + } +} + +fn report_lag(owner: &FpmTraceInner, dropped: u64, during_shutdown: bool) { + let total_dropped = owner.record_dropped(dropped); + tracing::warn!( + producer_id = %owner.producer_id, + dropped, + total_dropped, + during_shutdown, + "FPM trace queue lagged; older records were overwritten" + ); +} + +fn report_drop_summary(owner: &FpmTraceInner) { + let dropped = owner.dropped.load(Ordering::Relaxed); + if dropped > 0 { + tracing::warn!( + producer_id = %owner.producer_id, + dropped, + "FPM trace closed after dropping records" + ); + } +} + +async fn run_full( + mut receiver: broadcast::Receiver, + writer: JsonlGzipWriter, + owner: Arc, +) { + loop { + tokio::select! { + biased; + _ = owner.shutdown.cancelled() => { + // Stop new publishes and wait for publishers that crossed the + // acceptance boundary. The subsequent try_recv loop is then a + // stable drain of the bounded queue. + owner.stop_accepting().await; + loop { + match receiver.try_recv() { + Ok(event) => send_decoded(&writer, event).await, + Err(broadcast::error::TryRecvError::Lagged(dropped)) => { + report_lag(&owner, dropped, true); + } + Err( + broadcast::error::TryRecvError::Empty + | broadcast::error::TryRecvError::Closed + ) => break, + } + } + if let Err(error) = writer.close().await { + tracing::warn!(%error, "FPM trace writer failed to close cleanly"); + } + report_drop_summary(&owner); + return; + } + message = receiver.recv() => { + match message { + Ok(event) => send_decoded(&writer, event).await, + Err(broadcast::error::RecvError::Lagged(dropped)) => { + report_lag(&owner, dropped, false); + } + Err(broadcast::error::RecvError::Closed) => { + owner.stop_accepting().await; + if let Err(error) = writer.close().await { + tracing::warn!(%error, "FPM trace writer failed to close cleanly"); + } + report_drop_summary(&owner); + return; + } + } + } + } + } +} + +struct PendingSample { + counter_id: i64, + record: FpmTraceRecord, +} + +fn retain_latest( + latest: &mut BTreeMap, + last_emitted: &BTreeMap, + event: FpmTraceEvent, +) { + match decode_event(event, FpmTraceMode::Sampled) { + Ok((key, counter_id, record)) => { + if last_emitted.get(&key) == Some(&counter_id) { + return; + } + latest.insert(key, PendingSample { counter_id, record }); + } + Err(error) => tracing::warn!(%error, "FPM trace dropped malformed payload"), + } +} + +async fn flush_latest( + writer: &JsonlGzipWriter, + latest: &mut BTreeMap, + last_emitted: &mut BTreeMap, +) { + for (key, pending) in std::mem::take(latest) { + let counter_id = pending.counter_id; + if writer.send(pending.record).await.is_err() { + tracing::warn!("FPM trace writer closed; dropping sampled record"); + break; + } + last_emitted.insert(key, counter_id); + } +} + +async fn run_sampled( + mut receiver: broadcast::Receiver, + writer: JsonlGzipWriter, + owner: Arc, + sample_interval_ms: u64, +) { + let interval = Duration::from_millis(sample_interval_ms.max(1)); + let mut flush_tick = tokio::time::interval_at(tokio::time::Instant::now() + interval, interval); + flush_tick.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Delay); + let mut latest = BTreeMap::new(); + let mut last_emitted = BTreeMap::new(); + + loop { + tokio::select! { + biased; + _ = owner.shutdown.cancelled() => { + owner.stop_accepting().await; + loop { + match receiver.try_recv() { + Ok(event) => retain_latest(&mut latest, &last_emitted, event), + Err(broadcast::error::TryRecvError::Lagged(dropped)) => { + report_lag(&owner, dropped, true); + } + Err( + broadcast::error::TryRecvError::Empty + | broadcast::error::TryRecvError::Closed + ) => break, + } + } + flush_latest(&writer, &mut latest, &mut last_emitted).await; + if let Err(error) = writer.close().await { + tracing::warn!(%error, "FPM trace writer failed to close cleanly"); + } + report_drop_summary(&owner); + return; + } + _ = flush_tick.tick() => { + flush_latest(&writer, &mut latest, &mut last_emitted).await; + } + message = receiver.recv() => { + match message { + Ok(event) => retain_latest(&mut latest, &last_emitted, event), + Err(broadcast::error::RecvError::Lagged(dropped)) => { + report_lag(&owner, dropped, false); + } + Err(broadcast::error::RecvError::Closed) => { + owner.stop_accepting().await; + flush_latest(&writer, &mut latest, &mut last_emitted).await; + if let Err(error) = writer.close().await { + tracing::warn!(%error, "FPM trace writer failed to close cleanly"); + } + report_drop_summary(&owner); + return; + } + } + } + } + } +} + +#[cfg(test)] +mod tests { + use std::io::Read; + use std::sync::atomic::{AtomicBool, AtomicUsize}; + + use bytes::Bytes; + use flate2::read::MultiGzDecoder; + use tokio_util::sync::CancellationToken; + + use super::*; + use crate::telemetry::jsonl_gz::segment_path; + + fn source() -> FpmTraceSource { + FpmTraceSource { + namespace: "dynamo".to_string(), + component: "backend".to_string(), + producer_id: "producer-1".to_string(), + } + } + + fn event_with_wall_time( + worker_id: &str, + dp_rank: i64, + counter_id: i64, + wall_time: f64, + ) -> FpmTraceEvent { + let active = wall_time > 0.0; + let payload = rmp_serde::to_vec_named(&serde_json::json!({ + "version": 1, + "worker_id": worker_id, + "dp_rank": dp_rank, + "counter_id": counter_id, + "wall_time": wall_time, + "scheduled_requests": { + "num_prefill_requests": if active { 1 } else { 0 }, + "sum_prefill_tokens": if active { 16 } else { 0 }, + "var_prefill_length": 0.0, + "sum_prefill_kv_tokens": 0, + "num_decode_requests": if active { 2 } else { 0 }, + "sum_decode_kv_tokens": if active { 64 } else { 0 }, + "var_decode_kv_tokens": 0.0, + }, + "queued_requests": { + "num_prefill_requests": 0, + "sum_prefill_tokens": 0, + "var_prefill_length": 0.0, + "num_decode_requests": 0, + "sum_decode_kv_tokens": 0, + "var_decode_kv_tokens": 0.0, + }, + })) + .unwrap(); + FpmTraceEvent { + observed_at_unix_ms: counter_id as u64 * 100, + payload: Bytes::from(payload), + source: Arc::new(source()), + } + } + + fn event(worker_id: &str, dp_rank: i64, counter_id: i64) -> FpmTraceEvent { + event_with_wall_time(worker_id, dp_rank, counter_id, 0.01) + } + + fn key(worker_id: &str, dp_rank: i64) -> FpmKey { + (source(), worker_id.to_string(), dp_rank) + } + + async fn test_writer(path: &Path) -> JsonlGzipWriter { + JsonlGzipWriter::new( + path.display().to_string(), + JsonlGzipSinkOptions { + buffer_bytes: 1, + flush_interval: Duration::from_secs(60), + roll_uncompressed_bytes: 1024 * 1024, + roll_lines: None, + max_segments: Some(4), + }, + ) + .await + .unwrap() + } + + fn read_trace_records(path: &Path) -> Vec { + let bytes = std::fs::read(segment_path(path, 0)).unwrap(); + let mut decoder = MultiGzDecoder::new(bytes.as_slice()); + let mut content = String::new(); + decoder.read_to_string(&mut content).unwrap(); + content + .lines() + .map(|line| serde_json::from_str(line).unwrap()) + .collect() + } + + async fn yield_to_trace_worker() { + for _ in 0..8 { + tokio::task::yield_now().await; + } + } + + fn owner( + sender: broadcast::Sender, + shutdown: CancellationToken, + ) -> Arc { + Arc::new(FpmTraceInner { + producer_id: "producer-1".to_string(), + sender, + shutdown, + accepting: AtomicBool::new(true), + publishers_in_flight: AtomicUsize::new(0), + publishers_idle: tokio::sync::Notify::new(), + leases: AtomicUsize::new(1), + dropped: AtomicU64::new(0), + closed: AtomicBool::new(false), + closed_notify: tokio::sync::Notify::new(), + }) + } + + #[test] + fn producer_id_is_sanitized_and_inserted_before_segment_suffix() { + assert_eq!( + producer_output_path("/tmp/dynamo-fpm.jsonl.gz", "worker/a:b"), + "/tmp/dynamo-fpm.worker_a_b" + ); + assert_eq!( + producer_output_path("/tmp/dynamo-fpm", ""), + "/tmp/dynamo-fpm.unknown" + ); + } + + #[test] + fn writable_parent_preflight_rejects_a_blocked_path() { + let dir = tempfile::tempdir().unwrap(); + let blocker = dir.path().join("not-a-directory"); + std::fs::write(&blocker, b"file").unwrap(); + let output = blocker.join("dynamo-fpm"); + + assert!(preflight_writable_parent(&output.to_string_lossy()).is_err()); + } + + #[test] + fn writable_parent_preflight_removes_its_unique_probe() { + let dir = tempfile::tempdir().unwrap(); + let output = dir.path().join("dynamo-fpm.producer-7"); + + preflight_writable_parent(&output.to_string_lossy()).unwrap(); + + assert_eq!(std::fs::read_dir(dir.path()).unwrap().count(), 0); + } + + #[test] + fn sampled_mode_keeps_latest_event_per_worker_and_rank() { + let mut latest = BTreeMap::new(); + let last_emitted = BTreeMap::new(); + retain_latest(&mut latest, &last_emitted, event("worker-a", 0, 1)); + retain_latest(&mut latest, &last_emitted, event("worker-a", 0, 2)); + retain_latest(&mut latest, &last_emitted, event("worker-a", 1, 3)); + + assert_eq!(latest.len(), 2); + assert_eq!(latest[&key("worker-a", 0)].record.observed_at_unix_ms, 200); + assert_eq!(latest[&key("worker-a", 1)].record.observed_at_unix_ms, 300); + assert_eq!(latest[&key("worker-a", 0)].record.fpm["counter_id"], 2); + } + + #[test] + fn sampled_mode_keeps_components_with_overlapping_worker_ids_independent() { + let mut latest = BTreeMap::new(); + let last_emitted = BTreeMap::new(); + let backend = event("worker-a", 0, 1); + let mut prefill = event("worker-a", 0, 2); + prefill.source = Arc::new(FpmTraceSource { + component: "prefill".to_string(), + ..source() + }); + + retain_latest(&mut latest, &last_emitted, backend); + retain_latest(&mut latest, &last_emitted, prefill); + + assert_eq!(latest.len(), 2); + assert_eq!(latest[&key("worker-a", 0)].counter_id, 1); + let prefill_key = ( + FpmTraceSource { + component: "prefill".to_string(), + ..source() + }, + "worker-a".to_string(), + 0, + ); + assert_eq!(latest[&prefill_key].counter_id, 2); + } + + #[test] + fn sampled_state_stays_bounded_to_one_record_per_rank() { + let mut latest = BTreeMap::new(); + let last_emitted = BTreeMap::new(); + + for counter_id in 1..=32 { + for dp_rank in 0..256 { + retain_latest( + &mut latest, + &last_emitted, + event("worker-a", dp_rank, counter_id), + ); + } + } + + assert_eq!(latest.len(), 256); + assert!(latest.values().all(|pending| pending.counter_id == 32)); + } + + #[test] + fn sampled_mode_does_not_reemit_the_last_counter() { + let mut latest = BTreeMap::new(); + let mut last_emitted = BTreeMap::new(); + let key = key("worker-a", 0); + last_emitted.insert(key, 9); + + retain_latest(&mut latest, &last_emitted, event("worker-a", 0, 9)); + + assert!(latest.is_empty()); + } + + #[test] + fn malformed_payload_is_rejected() { + let malformed = FpmTraceEvent { + observed_at_unix_ms: 1, + payload: Bytes::from_static(b"not-msgpack"), + source: Arc::new(source()), + }; + assert!(decode_event(malformed, FpmTraceMode::Full).is_err()); + } + + #[test] + fn incomplete_fpm_map_is_rejected() { + let payload = rmp_serde::to_vec_named(&serde_json::json!({ + "version": 1, + "worker_id": "worker-a", + "dp_rank": 0, + "counter_id": 1, + "wall_time": 0.01, + "scheduled_requests": {}, + "queued_requests": {}, + })) + .unwrap(); + let incomplete = FpmTraceEvent { + observed_at_unix_ms: 1, + payload: Bytes::from(payload), + source: Arc::new(source()), + }; + assert!(decode_event(incomplete, FpmTraceMode::Full).is_err()); + } + + #[test] + fn trace_record_carries_stable_schema_source_and_mode() { + let (_, _, record) = decode_event(event("worker-a", 2, 7), FpmTraceMode::Sampled).unwrap(); + let value = serde_json::to_value(record).unwrap(); + assert_eq!(value["schema"], "dynamo.fpm.trace.v1"); + assert_eq!(value["source"]["namespace"], "dynamo"); + assert_eq!(value["source"]["component"], "backend"); + assert_eq!(value["source"]["producer_id"], "producer-1"); + assert_eq!(value["capture_mode"], "sampled"); + assert_eq!(value["fpm"]["counter_id"], 7); + } + + #[tokio::test] + async fn full_mode_persists_active_and_idle_payloads_in_order() { + let dir = tempfile::tempdir().unwrap(); + let path = dir.path().join("full-trace"); + let writer = test_writer(&path).await; + let (sender, receiver) = broadcast::channel(8); + let shutdown = CancellationToken::new(); + let owner = owner(sender.clone(), shutdown.clone()); + let task = tokio::spawn(run_full(receiver, writer, owner)); + + sender + .send(event_with_wall_time("worker-a", 0, 1, 0.01)) + .unwrap(); + sender + .send(event_with_wall_time("worker-a", 0, 2, 0.0)) + .unwrap(); + sender + .send(event_with_wall_time("worker-a", 0, 3, 0.02)) + .unwrap(); + shutdown.cancel(); + task.await.unwrap(); + + let records = read_trace_records(&path); + let counters: Vec<_> = records + .iter() + .map(|record| record["event"]["fpm"]["counter_id"].as_i64().unwrap()) + .collect(); + assert_eq!(counters, [1, 2, 3]); + assert_eq!(records[1]["event"]["fpm"]["wall_time"], 0.0); + assert!( + records + .iter() + .all(|record| record["event"]["capture_mode"] == "full") + ); + } + + #[tokio::test(start_paused = true)] + async fn sampled_mode_is_per_rank_suppresses_unchanged_and_flushes_dirty_shutdown() { + let dir = tempfile::tempdir().unwrap(); + let path = dir.path().join("sampled-trace"); + let writer = test_writer(&path).await; + let (sender, receiver) = broadcast::channel(16); + let shutdown = CancellationToken::new(); + let owner = owner(sender.clone(), shutdown.clone()); + let task = tokio::spawn(run_sampled(receiver, writer, owner, 100)); + + sender.send(event("worker-a", 0, 1)).unwrap(); + sender.send(event("worker-a", 0, 2)).unwrap(); + sender.send(event("worker-a", 1, 1)).unwrap(); + yield_to_trace_worker().await; + tokio::time::advance(Duration::from_millis(100)).await; + yield_to_trace_worker().await; + + // Repeating the last emitted counters in a later interval must not + // create duplicate samples. + sender.send(event("worker-a", 0, 2)).unwrap(); + sender.send(event("worker-a", 1, 1)).unwrap(); + yield_to_trace_worker().await; + tokio::time::advance(Duration::from_millis(100)).await; + yield_to_trace_worker().await; + + // This sample is dirty but its interval has not elapsed; cancellation + // must still flush it and await the gzip writer close. + sender.send(event("worker-a", 0, 3)).unwrap(); + shutdown.cancel(); + task.await.unwrap(); + + let records = read_trace_records(&path); + let rank_counters: Vec<_> = records + .iter() + .map(|record| { + ( + record["event"]["fpm"]["dp_rank"].as_i64().unwrap(), + record["event"]["fpm"]["counter_id"].as_i64().unwrap(), + ) + }) + .collect(); + assert_eq!(rank_counters, [(0, 2), (1, 1), (0, 3)]); + assert!( + records + .iter() + .all(|record| record["event"]["capture_mode"] == "sampled") + ); + } + + #[tokio::test] + async fn bounded_trace_queue_retains_newest_and_counts_overwritten_records() { + let (sender, mut receiver) = broadcast::channel(2); + let owner = owner(sender.clone(), CancellationToken::new()); + sender.send(event("worker-a", 0, 1)).unwrap(); + sender.send(event("worker-a", 0, 2)).unwrap(); + sender.send(event("worker-a", 0, 3)).unwrap(); + + let dropped = match receiver.recv().await { + Err(broadcast::error::RecvError::Lagged(dropped)) => dropped, + result => panic!("expected lag, got {result:?}"), + }; + report_lag(&owner, dropped, false); + assert_eq!(owner.dropped.load(Ordering::Relaxed), 1); + + let second = receiver.recv().await.unwrap(); + let third = receiver.recv().await.unwrap(); + let (_, second_counter, _) = decode_event(second, FpmTraceMode::Full).unwrap(); + let (_, third_counter, _) = decode_event(third, FpmTraceMode::Full).unwrap(); + assert_eq!((second_counter, third_counter), (2, 3)); + } + + #[tokio::test] + async fn worker_completion_notifies_registry_even_after_task_panic() { + let (sender, _receiver) = broadcast::channel(1); + let owner = owner(sender, CancellationToken::new()); + let completion = MarkClosedOnDrop(owner.clone()); + + let task = tokio::spawn(async move { + let _completion = completion; + panic!("simulated trace worker panic"); + }); + assert!(task.await.unwrap_err().is_panic()); + owner.wait_closed().await; + assert!(owner.closed.load(Ordering::Acquire)); + assert!(!owner.accepting.load(Ordering::Acquire)); + } +} diff --git a/lib/llm/src/lib.rs b/lib/llm/src/lib.rs index 27083abc768b..b85a1ce7e2c6 100644 --- a/lib/llm/src/lib.rs +++ b/lib/llm/src/lib.rs @@ -13,6 +13,7 @@ pub mod endpoint_type; pub mod engines; pub mod entrypoint; pub mod fpm_publisher; +pub mod fpm_trace; pub mod frontend_config; pub mod grpc; pub mod http; diff --git a/lib/llm/src/request_trace/sink.rs b/lib/llm/src/request_trace/sink.rs index a3c9049843db..53295532a95d 100644 --- a/lib/llm/src/request_trace/sink.rs +++ b/lib/llm/src/request_trace/sink.rs @@ -117,6 +117,7 @@ impl JsonlGzipRequestTraceSink { flush_interval: Duration::from_millis(policy.jsonl_flush_interval_ms.max(1)), roll_uncompressed_bytes: policy.jsonl_gz_roll_bytes, roll_lines: policy.jsonl_gz_roll_lines, + max_segments: None, }, ) .await @@ -303,6 +304,7 @@ mod tests { flush_interval: Duration::from_secs(60), roll_uncompressed_bytes: 1024 * 1024, roll_lines: Some(1), + max_segments: None, }, ) .await diff --git a/lib/llm/src/telemetry/jsonl_gz.rs b/lib/llm/src/telemetry/jsonl_gz.rs index 9945ffea86e3..aa46d5ac5087 100644 --- a/lib/llm/src/telemetry/jsonl_gz.rs +++ b/lib/llm/src/telemetry/jsonl_gz.rs @@ -7,6 +7,7 @@ //! batched into uncompressed bytes, and appended as gzip members. Segments roll //! when uncompressed bytes or record-line thresholds are exceeded. +use std::fs::File; use std::io::{BufWriter, Write}; use std::path::{Path, PathBuf}; use std::time::{Duration, Instant}; @@ -23,6 +24,9 @@ pub struct JsonlGzipSinkOptions { pub flush_interval: Duration, pub roll_uncompressed_bytes: u64, pub roll_lines: Option, + /// Maximum number of segments to retain for this exact output prefix. + /// `None` preserves all segments. + pub max_segments: Option, } impl Default for JsonlGzipSinkOptions { @@ -32,15 +36,19 @@ impl Default for JsonlGzipSinkOptions { flush_interval: Duration::from_millis(1000), roll_uncompressed_bytes: 256 * 1024 * 1024, roll_lines: None, + max_segments: None, } } } -/// Channel-backed handle for a rotating gzip JSONL sink. Drop cancels the -/// writer task; remaining records are flushed before exit. +/// Channel-backed handle for a rotating gzip JSONL sink. +/// +/// Drop requests writer shutdown, but cannot wait for its final flush. Call +/// [`Self::shutdown`] or [`Self::close`] when completion must be awaited. pub struct JsonlGzipWriter { - tx: mpsc::Sender, + tx: Option>, shutdown: CancellationToken, + worker: Option>>, } #[derive(Serialize)] @@ -56,21 +64,47 @@ where pub async fn new(path: String, options: JsonlGzipSinkOptions) -> anyhow::Result { let shutdown = CancellationToken::new(); let (tx, rx) = mpsc::channel::(2048); - let mut writer = GzipBatchWriter::new(path.clone(), options) - .with_context(|| format!("opening gzip jsonl sink at {path}"))?; + let display_path = path.clone(); + let mut writer = + tokio::task::spawn_blocking(move || GzipBatchWriter::::new(path, options)) + .await + .context("gzip jsonl sink initializer panicked")? + .with_context(|| format!("opening gzip jsonl sink at {display_path}"))?; let worker_shutdown = shutdown.clone(); - tokio::spawn(async move { - run_gzip_writer(rx, &mut writer, worker_shutdown).await; - }); + let worker = + tokio::spawn(async move { run_gzip_writer(rx, &mut writer, worker_shutdown).await }); - Ok(Self { tx, shutdown }) + Ok(Self { + tx: Some(tx), + shutdown, + worker: Some(worker), + }) } /// Forward a record to the writer task. Returns `Err` if the worker has /// shut down. pub async fn send(&self, rec: T) -> Result<(), mpsc::error::SendError> { - self.tx.send(rec).await + match &self.tx { + Some(tx) => tx.send(rec).await, + None => Err(mpsc::error::SendError(rec)), + } + } + + /// Drain all accepted records, flush the active segment, and wait for the + /// writer task to exit. + pub async fn shutdown(&mut self) -> anyhow::Result<()> { + self.tx.take(); + self.shutdown.cancel(); + if let Some(worker) = self.worker.take() { + worker.await.context("gzip jsonl writer task panicked")??; + } + Ok(()) + } + + /// Consuming convenience wrapper around [`Self::shutdown`]. + pub async fn close(mut self) -> anyhow::Result<()> { + self.shutdown().await } } @@ -85,6 +119,8 @@ struct GzipBatchWriter { current_index: u64, start_time: Instant, batch: Vec, + active_file: Option, + prune_pending: bool, segment_uncompressed_bytes: u64, segment_lines: u64, options: JsonlGzipSinkOptions, @@ -93,6 +129,10 @@ struct GzipBatchWriter { impl GzipBatchWriter { fn new(path: String, options: JsonlGzipSinkOptions) -> anyhow::Result { + if options.max_segments == Some(0) { + return Err(anyhow!("gzip jsonl max_segments must be positive")); + } + let base_path = PathBuf::from(path); if let Some(parent) = base_path.parent() && !parent.as_os_str().is_empty() @@ -107,6 +147,8 @@ impl GzipBatchWriter { current_index, start_time: Instant::now(), batch: Vec::with_capacity(options.buffer_bytes.max(1)), + active_file: None, + prune_pending: true, segment_uncompressed_bytes: 0, segment_lines: 0, options, @@ -159,6 +201,8 @@ impl GzipBatchWriter { fn roll_segment(&mut self) { self.current_index = self.current_index.saturating_add(1); + self.active_file = None; + self.prune_pending = true; self.segment_uncompressed_bytes = 0; self.segment_lines = 0; } @@ -168,12 +212,41 @@ impl GzipBatchWriter { return Ok(()); } - let path = segment_path(&self.base_path, self.current_index); let batch = std::mem::take(&mut self.batch); + let base_path = self.base_path.clone(); + let current_index = self.current_index; + let max_segments = self.options.max_segments; + let active_file = self.active_file.take(); + let prune_pending = self.prune_pending; + + let (active_file, current_index, result) = tokio::task::spawn_blocking(move || { + let (mut active_file, current_index, path) = match active_file { + Some(file) => (file, current_index, segment_path(&base_path, current_index)), + None => match create_available_segment(&base_path, current_index) { + Ok(segment) => segment, + Err(err) => return (None, current_index, Err(err)), + }, + }; + + let result = write_gzip_member(&mut active_file, &path, batch).map(|()| { + if prune_pending + && let Some(max_segments) = max_segments + && let Err(err) = prune_segments(&base_path, max_segments) + { + tracing::warn!("gzip jsonl sink failed to prune old segments: {err}"); + } + }); + (Some(active_file), current_index, result) + }) + .await + .context("gzip jsonl writer task panicked")?; - tokio::task::spawn_blocking(move || write_gzip_member(path, batch)) - .await - .context("gzip jsonl writer task panicked")??; + self.active_file = active_file; + self.current_index = current_index; + if result.is_ok() { + self.prune_pending = false; + } + result?; Ok(()) } @@ -183,10 +256,11 @@ async fn run_gzip_writer( mut rx: mpsc::Receiver, writer: &mut GzipBatchWriter, shutdown: CancellationToken, -) { +) -> anyhow::Result<()> { let mut flush_tick = tokio::time::interval(writer.options.flush_interval.max(Duration::from_millis(1))); flush_tick.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Delay); + let mut first_error = None; loop { tokio::select! { @@ -195,16 +269,19 @@ async fn run_gzip_writer( while let Ok(rec) = rx.try_recv() { if let Err(err) = writer.push(&rec).await { tracing::warn!("gzip jsonl sink dropped record during shutdown: {err}"); + first_error.get_or_insert(err); } } if let Err(err) = writer.flush_batch().await { tracing::warn!("gzip jsonl sink failed final flush: {err}"); + first_error.get_or_insert(err); } - return; + return first_error.map_or(Ok(()), Err); } _ = flush_tick.tick() => { if let Err(err) = writer.flush_batch().await { tracing::warn!("gzip jsonl sink failed flush: {err}"); + first_error.get_or_insert(err); } } msg = rx.recv() => { @@ -212,13 +289,15 @@ async fn run_gzip_writer( Some(rec) => { if let Err(err) = writer.push(&rec).await { tracing::warn!("gzip jsonl sink dropped record: {err}"); + first_error.get_or_insert(err); } } None => { if let Err(err) = writer.flush_batch().await { tracing::warn!("gzip jsonl sink failed final flush: {err}"); + first_error.get_or_insert(err); } - return; + return first_error.map_or(Ok(()), Err); } } } @@ -226,19 +305,55 @@ async fn run_gzip_writer( } } -fn write_gzip_member(path: PathBuf, batch: Vec) -> anyhow::Result<()> { +fn ensure_segment_parent(path: &Path) -> anyhow::Result<()> { if let Some(parent) = path.parent() && !parent.as_os_str().is_empty() { std::fs::create_dir_all(parent) .with_context(|| format!("creating gzip jsonl directory {}", parent.display()))?; } + Ok(()) +} - let file = std::fs::OpenOptions::new() - .create(true) - .append(true) - .open(&path) - .with_context(|| format!("opening gzip jsonl segment {}", path.display()))?; +fn create_available_segment( + base_path: &Path, + initial_index: u64, +) -> anyhow::Result<(File, u64, PathBuf)> { + let initial_path = segment_path(base_path, initial_index); + // Check the parent outside the collision loop. An `AlreadyExists` error + // from directory creation is not evidence that a segment index is taken. + ensure_segment_parent(&initial_path)?; + + let mut index = initial_index; + loop { + let path = segment_path(base_path, index); + // `create_new` is atomic and fails for every existing filesystem + // object, including symlinks. Keeping this handle for the segment's + // lifetime also prevents later flushes from following a replacement + // symlink. + match std::fs::OpenOptions::new() + .create_new(true) + .append(true) + .open(&path) + { + Ok(file) => return Ok((file, index, path)), + Err(error) if error.kind() == std::io::ErrorKind::AlreadyExists => { + index = index.checked_add(1).ok_or_else(|| { + anyhow!( + "no available gzip jsonl segment index for {}", + base_path.display() + ) + })?; + } + Err(error) => { + return Err(error) + .with_context(|| format!("creating gzip jsonl segment {}", path.display())); + } + } + } +} + +fn write_gzip_member(file: &mut File, path: &Path, batch: Vec) -> anyhow::Result<()> { let writer = BufWriter::new(file); let mut encoder = GzEncoder::new(writer, Compression::default()); encoder @@ -263,15 +378,89 @@ pub fn segment_path(base_path: &Path, index: u64) -> PathBuf { } fn next_segment_index(base_path: &Path) -> anyhow::Result { - for index in 0..u64::MAX { - if !segment_path(base_path, index).try_exists()? { - return Ok(index); + match matching_segments(base_path)?.occupied.last() { + Some((index, _)) => index.checked_add(1).ok_or_else(|| { + anyhow!( + "no available gzip jsonl segment index for {}", + base_path.display() + ) + }), + None => Ok(0), + } +} + +fn prune_segments(base_path: &Path, max_segments: usize) -> anyhow::Result<()> { + prune_segments_with(base_path, max_segments, |path| { + std::fs::remove_file(path) + .with_context(|| format!("pruning gzip jsonl segment {}", path.display())) + }) +} + +fn prune_segments_with( + base_path: &Path, + max_segments: usize, + mut remove: impl FnMut(&Path) -> anyhow::Result<()>, +) -> anyhow::Result<()> { + let segments = matching_segments(base_path)?.prunable; + let remove_count = segments.len().saturating_sub(max_segments); + for (_, path) in segments.into_iter().take(remove_count) { + remove(&path)?; + } + Ok(()) +} + +struct MatchingSegments { + occupied: Vec<(u64, PathBuf)>, + prunable: Vec<(u64, PathBuf)>, +} + +fn matching_segments(base_path: &Path) -> anyhow::Result { + let first_segment = segment_path(base_path, 0); + let parent = first_segment + .parent() + .filter(|path| !path.as_os_str().is_empty()) + .unwrap_or_else(|| Path::new(".")); + if !parent.try_exists()? { + return Ok(MatchingSegments { + occupied: Vec::new(), + prunable: Vec::new(), + }); + } + + let mut occupied = Vec::new(); + let mut prunable = Vec::new(); + for entry in std::fs::read_dir(parent) + .with_context(|| format!("listing gzip jsonl directory {}", parent.display()))? + { + let entry = entry.with_context(|| { + format!("reading gzip jsonl directory entry in {}", parent.display()) + })?; + let Some(index) = segment_index(base_path, &entry.path()) else { + continue; + }; + let path = entry.path(); + let file_type = entry + .file_type() + .with_context(|| format!("reading file type for {}", path.display()))?; + occupied.push((index, path.clone())); + if file_type.is_file() { + prunable.push((index, path)); } } - Err(anyhow!( - "no available gzip jsonl segment index for {}", - base_path.display() - )) + occupied.sort_unstable_by_key(|(index, _)| *index); + prunable.sort_unstable_by_key(|(index, _)| *index); + Ok(MatchingSegments { occupied, prunable }) +} + +fn segment_index(base_path: &Path, candidate: &Path) -> Option { + let name = candidate.file_name()?.to_str()?; + let without_suffix = name.strip_suffix(".jsonl.gz")?; + let (_, index_text) = without_suffix.rsplit_once('.')?; + if index_text.len() < 6 || !index_text.bytes().all(|byte| byte.is_ascii_digit()) { + return None; + } + let index = index_text.parse::().ok()?; + (segment_path(base_path, index).file_name()? == candidate.file_name()?).then_some(index) } #[cfg(test)] @@ -311,6 +500,7 @@ mod tests { flush_interval: Duration::from_secs(60), roll_uncompressed_bytes: 1024 * 1024, roll_lines: None, + max_segments: None, }, ) .await @@ -330,18 +520,10 @@ mod tests { }) .await .unwrap(); + writer.close().await.expect("writer should close cleanly"); let segment = segment_path(&path, 0); - let mut content = String::new(); - for _ in 0..100 { - if segment.exists() { - content = read_gzip_jsonl(&segment); - if content.matches("\"name\":").count() == 2 { - break; - } - } - tokio::time::sleep(Duration::from_millis(20)).await; - } + let content = read_gzip_jsonl(&segment); assert!(content.contains("\"name\":\"first\"")); assert!(content.contains("\"name\":\"second\"")); } @@ -357,6 +539,7 @@ mod tests { flush_interval: Duration::from_secs(60), roll_uncompressed_bytes: 1024 * 1024, roll_lines: Some(1), + max_segments: None, }, ) .await @@ -376,25 +559,415 @@ mod tests { }) .await .unwrap(); + writer.close().await.expect("writer should close cleanly"); let first = segment_path(&path, 0); let second = segment_path(&path, 1); - let mut first_content = String::new(); - let mut second_content = String::new(); - for _ in 0..100 { - if first.exists() && second.exists() { - first_content = read_gzip_jsonl(&first); - second_content = read_gzip_jsonl(&second); - if first_content.contains("\"name\":\"first\"") - && second_content.contains("\"name\":\"second\"") - { - break; - } - } - tokio::time::sleep(Duration::from_millis(20)).await; - } + let first_content = read_gzip_jsonl(&first); + let second_content = read_gzip_jsonl(&second); assert!(first_content.contains("\"name\":\"first\"")); assert!(second_content.contains("\"name\":\"second\"")); } + + #[tokio::test] + async fn restart_uses_one_more_than_highest_existing_index() { + let dir = tempdir().unwrap(); + let path = dir.path().join("test_trace"); + std::fs::write(segment_path(&path, 0), b"old zero").unwrap(); + std::fs::write(segment_path(&path, 2), b"old two").unwrap(); + + let writer: JsonlGzipWriter = JsonlGzipWriter::new( + path.display().to_string(), + JsonlGzipSinkOptions { + buffer_bytes: 1024, + flush_interval: Duration::from_secs(60), + roll_uncompressed_bytes: 1024 * 1024, + roll_lines: None, + max_segments: None, + }, + ) + .await + .unwrap(); + writer + .send(TestRecord { + id: 3, + name: "after restart".to_string(), + }) + .await + .unwrap(); + writer.close().await.unwrap(); + + assert!(!segment_path(&path, 1).exists()); + assert!(segment_path(&path, 0).exists()); + assert!(segment_path(&path, 2).exists()); + assert!(segment_path(&path, 3).exists()); + assert!(read_gzip_jsonl(&segment_path(&path, 3)).contains("after restart")); + } + + #[tokio::test] + async fn overlapping_writers_claim_distinct_segments_on_first_flush() { + let dir = tempdir().unwrap(); + let path = dir.path().join("test_trace"); + let options = JsonlGzipSinkOptions { + buffer_bytes: 1024, + flush_interval: Duration::from_secs(60), + roll_uncompressed_bytes: 1024 * 1024, + roll_lines: None, + max_segments: None, + }; + let first = JsonlGzipWriter::::new(path.display().to_string(), options.clone()) + .await + .unwrap(); + let second = JsonlGzipWriter::::new(path.display().to_string(), options) + .await + .unwrap(); + + first + .send(TestRecord { + id: 1, + name: "first writer".to_string(), + }) + .await + .unwrap(); + first.close().await.unwrap(); + second + .send(TestRecord { + id: 2, + name: "second writer".to_string(), + }) + .await + .unwrap(); + second.close().await.unwrap(); + + assert!(read_gzip_jsonl(&segment_path(&path, 0)).contains("first writer")); + assert!(read_gzip_jsonl(&segment_path(&path, 1)).contains("second writer")); + } + + #[tokio::test] + async fn retention_prunes_only_exact_prefix_after_new_segment_is_written() { + let dir = tempdir().unwrap(); + let path = dir.path().join("test_trace"); + let unrelated = dir.path().join("test_trace_other.000000.jsonl.gz"); + let backup = dir.path().join("test_trace.000000.jsonl.gz.bak"); + std::fs::write(segment_path(&path, 0), b"old zero").unwrap(); + std::fs::write(segment_path(&path, 1), b"old one").unwrap(); + std::fs::write(&unrelated, b"unrelated").unwrap(); + std::fs::write(&backup, b"backup").unwrap(); + + let writer: JsonlGzipWriter = JsonlGzipWriter::new( + path.display().to_string(), + JsonlGzipSinkOptions { + buffer_bytes: 1024, + flush_interval: Duration::from_secs(60), + roll_uncompressed_bytes: 1024 * 1024, + roll_lines: None, + max_segments: Some(2), + }, + ) + .await + .unwrap(); + writer + .send(TestRecord { + id: 2, + name: "replacement".to_string(), + }) + .await + .unwrap(); + writer.close().await.unwrap(); + + assert!(!segment_path(&path, 0).exists()); + assert!(segment_path(&path, 1).exists()); + assert!(segment_path(&path, 2).exists()); + assert!(unrelated.exists()); + assert!(backup.exists()); + } + + #[tokio::test] + async fn repeated_flush_of_active_segment_does_not_rescan_retention() { + let dir = tempdir().unwrap(); + let path = dir.path().join("test_trace"); + let oldest = segment_path(&path, 0); + std::fs::write(&oldest, b"old zero").unwrap(); + std::fs::write(segment_path(&path, 1), b"old one").unwrap(); + + let mut writer = GzipBatchWriter::::new( + path.display().to_string(), + JsonlGzipSinkOptions { + buffer_bytes: 1024, + flush_interval: Duration::from_secs(60), + roll_uncompressed_bytes: 1024 * 1024, + roll_lines: None, + max_segments: Some(2), + }, + ) + .unwrap(); + writer + .push(&TestRecord { + id: 2, + name: "first flush".to_string(), + }) + .await + .unwrap(); + writer.flush_batch().await.unwrap(); + + assert!(!oldest.exists()); + std::fs::write(&oldest, b"created after the segment was opened").unwrap(); + + writer + .push(&TestRecord { + id: 3, + name: "second flush".to_string(), + }) + .await + .unwrap(); + writer.flush_batch().await.unwrap(); + + assert!(oldest.exists()); + let content = read_gzip_jsonl(&segment_path(&path, 2)); + assert!(content.contains("first flush")); + assert!(content.contains("second flush")); + } + + #[tokio::test] + async fn retention_of_one_keeps_only_active_segment() { + let dir = tempdir().unwrap(); + let path = dir.path().join("test_trace"); + let writer: JsonlGzipWriter = JsonlGzipWriter::new( + path.display().to_string(), + JsonlGzipSinkOptions { + buffer_bytes: 1, + flush_interval: Duration::from_secs(60), + roll_uncompressed_bytes: 1024 * 1024, + roll_lines: Some(1), + max_segments: Some(1), + }, + ) + .await + .unwrap(); + + for id in 0..3 { + writer + .send(TestRecord { + id, + name: format!("record {id}"), + }) + .await + .unwrap(); + } + writer.close().await.unwrap(); + + assert!(!segment_path(&path, 0).exists()); + assert!(!segment_path(&path, 1).exists()); + assert!(segment_path(&path, 2).exists()); + assert!(read_gzip_jsonl(&segment_path(&path, 2)).contains("record 2")); + } + + #[tokio::test] + async fn retention_keeps_four_newest_segments() { + let dir = tempdir().unwrap(); + let path = dir.path().join("test_trace"); + let writer: JsonlGzipWriter = JsonlGzipWriter::new( + path.display().to_string(), + JsonlGzipSinkOptions { + buffer_bytes: 1, + flush_interval: Duration::from_secs(60), + roll_uncompressed_bytes: 1024 * 1024, + roll_lines: Some(1), + max_segments: Some(4), + }, + ) + .await + .unwrap(); + + for id in 0..6 { + writer + .send(TestRecord { + id, + name: format!("record {id}"), + }) + .await + .unwrap(); + } + writer.close().await.unwrap(); + + assert!(!segment_path(&path, 0).exists()); + assert!(!segment_path(&path, 1).exists()); + for index in 2..6 { + assert!(segment_path(&path, index).exists()); + } + } + + #[tokio::test] + async fn oversized_record_occupies_a_segment_by_itself() { + let dir = tempdir().unwrap(); + let path = dir.path().join("test_trace"); + let writer: JsonlGzipWriter = JsonlGzipWriter::new( + path.display().to_string(), + JsonlGzipSinkOptions { + buffer_bytes: 1024 * 1024, + flush_interval: Duration::from_secs(60), + roll_uncompressed_bytes: 32, + roll_lines: None, + max_segments: None, + }, + ) + .await + .unwrap(); + writer + .send(TestRecord { + id: 0, + name: "x".repeat(1024), + }) + .await + .unwrap(); + writer + .send(TestRecord { + id: 1, + name: "next".to_string(), + }) + .await + .unwrap(); + writer.close().await.unwrap(); + + let first = read_gzip_jsonl(&segment_path(&path, 0)); + let second = read_gzip_jsonl(&segment_path(&path, 1)); + assert_eq!(first.matches("\"name\":").count(), 1); + assert!(first.contains(&"x".repeat(1024))); + assert_eq!(second.matches("\"name\":").count(), 1); + assert!(second.contains("\"name\":\"next\"")); + } + + #[tokio::test] + async fn occupied_directory_uses_next_index_and_is_not_pruned() { + let dir = tempdir().unwrap(); + let path = dir.path().join("test_trace"); + let old = segment_path(&path, 0); + let occupied = segment_path(&path, 1); + std::fs::write(&old, b"old").unwrap(); + std::fs::create_dir(&occupied).unwrap(); + + let writer: JsonlGzipWriter = JsonlGzipWriter::new( + path.display().to_string(), + JsonlGzipSinkOptions { + buffer_bytes: 1024, + flush_interval: Duration::from_secs(60), + roll_uncompressed_bytes: 1024 * 1024, + roll_lines: None, + max_segments: Some(1), + }, + ) + .await + .unwrap(); + writer + .send(TestRecord { + id: 1, + name: "cannot write".to_string(), + }) + .await + .unwrap(); + writer.close().await.unwrap(); + + assert!(!old.exists()); + assert!(occupied.is_dir()); + assert!(read_gzip_jsonl(&segment_path(&path, 2)).contains("cannot write")); + } + + #[cfg(unix)] + #[tokio::test] + async fn segment_creation_skips_symlink_created_after_allocation() { + use std::os::unix::fs::symlink; + + let dir = tempdir().unwrap(); + let path = dir.path().join("test_trace"); + let target = dir.path().join("target"); + let original = b"must remain unchanged"; + std::fs::write(&target, original).unwrap(); + + let mut writer = GzipBatchWriter::::new( + path.display().to_string(), + JsonlGzipSinkOptions { + buffer_bytes: 1024, + flush_interval: Duration::from_secs(60), + roll_uncompressed_bytes: 1024 * 1024, + roll_lines: None, + max_segments: Some(1), + }, + ) + .unwrap(); + symlink(&target, segment_path(&path, 0)).unwrap(); + + writer + .push(&TestRecord { + id: 1, + name: "must not reach target".to_string(), + }) + .await + .unwrap(); + writer.flush_batch().await.unwrap(); + + assert_eq!(std::fs::read(&target).unwrap(), original); + assert!( + std::fs::symlink_metadata(segment_path(&path, 0)) + .unwrap() + .file_type() + .is_symlink() + ); + assert!(read_gzip_jsonl(&segment_path(&path, 1)).contains("must not reach target")); + } + + #[test] + fn rolls_only_after_uncompressed_limit_would_be_exceeded() { + let dir = tempdir().unwrap(); + let path = dir.path().join("test_trace"); + let mut writer = GzipBatchWriter::::new( + path.display().to_string(), + JsonlGzipSinkOptions { + roll_uncompressed_bytes: 100, + ..Default::default() + }, + ) + .unwrap(); + writer.segment_lines = 1; + writer.segment_uncompressed_bytes = 50; + + assert!(!writer.should_roll_before(50)); + assert!(writer.should_roll_before(51)); + } + + #[test] + fn prune_failure_leaves_new_segment_and_returns_error() { + let dir = tempdir().unwrap(); + let path = dir.path().join("test_trace"); + for index in 0..3 { + std::fs::write(segment_path(&path, index), format!("segment {index}")).unwrap(); + } + + let mut attempted = Vec::new(); + let result = prune_segments_with(&path, 2, |candidate| { + attempted.push(candidate.to_path_buf()); + Err(anyhow!("injected prune failure")) + }); + + assert!(result.is_err()); + assert_eq!(attempted, vec![segment_path(&path, 0)]); + assert!(segment_path(&path, 0).exists()); + assert!(segment_path(&path, 2).exists()); + } + + #[tokio::test] + async fn rejects_zero_segment_retention() { + let dir = tempdir().unwrap(); + let path = dir.path().join("test_trace"); + let result = JsonlGzipWriter::::new( + path.display().to_string(), + JsonlGzipSinkOptions { + max_segments: Some(0), + ..Default::default() + }, + ) + .await; + + assert!(result.is_err()); + } } diff --git a/lib/runtime/src/config/environment_names.rs b/lib/runtime/src/config/environment_names.rs index 48425f08ce1b..5b6b32223d68 100644 --- a/lib/runtime/src/config/environment_names.rs +++ b/lib/runtime/src/config/environment_names.rs @@ -398,6 +398,30 @@ pub mod llm { pub const HISTOGRAM_PREFIX: &str = "DYN_HISTOGRAM_"; } + /// Forward-pass-metrics trace configuration. + pub mod fpm_trace { + /// Master switch. Truthy values persist locally produced FPM events. + pub const DYN_FPM_TRACE: &str = "DYN_FPM_TRACE"; + + /// Local gzip JSONL segment prefix. A sanitized producer id is appended + /// before the segment index so multiple producers do not share files. + pub const DYN_FPM_OUTPUT_PATH: &str = "DYN_FPM_OUTPUT_PATH"; + + /// Capture mode: `sampled` (latest event per DP rank each interval) or + /// `full` (every event reaching the producer-side trace tap). + pub const DYN_FPM_MODE: &str = "DYN_FPM_MODE"; + + /// Sampling interval in milliseconds when `DYN_FPM_MODE=sampled`. + pub const DYN_FPM_SAMPLE_INTERVAL_MS: &str = "DYN_FPM_SAMPLE_INTERVAL_MS"; + + /// Rotating gzip JSONL threshold in uncompressed bytes. + pub const DYN_FPM_JSONL_GZ_ROLL_BYTES: &str = "DYN_FPM_JSONL_GZ_ROLL_BYTES"; + + /// Maximum number of gzip JSONL segments retained per producer, + /// including the active segment. + pub const DYN_FPM_MAX_SEGMENTS: &str = "DYN_FPM_MAX_SEGMENTS"; + } + /// Audit sink configuration pub mod audit { /// Audit sink selection. Comma-separated values: `stderr`, `nats`,