feat: add opt-in memray allocation tracking (agent, manager, storage-proxy, app-proxy)

configs/agent/sample.toml

@@ -468,6 +468,29 @@
   # Added in 24.12.1
   ## sample-rate = 100
 
+# Memray allocation-tracking configuration for the agent. Enable it only while
+# diagnosing agent memory growth: the capture file keeps growing for as long as
+# the process runs.
+# Added in 26.8.0
+[memray]
+  # Whether to track memory allocations with memray. The capture file grows
+  # continuously while the process runs, so keep this disabled unless you are
+  # actively debugging memory usage.
+  # Added in 26.8.0
+  enabled = false
+  # Path to store the memray allocation capture. Memray refuses to start when
+  # the file already exists, so the process appends its PID to the path. Ensure
+  # the parent directory exists and has enough space.
+  # Added in 26.8.0
+  output-destination = "/var/log/backend.ai/memray/agent.bin"
+  # Whether to also capture native (C/C++) call stacks, so that allocations made
+  # inside extension modules are attributed to their native frames instead of
+  # the calling Python frame. This adds runtime overhead and enlarges the
+  # capture, and the native frames only resolve when debug symbols are
+  # available.
+  # Added in 26.8.0
+  native-traces = false
+
 # Logging configuration for the agent. Controls log levels, output destinations,
 # and log formatting. Proper logging setup is essential for debugging and
 # monitoring agent behavior in production environments.

src/ai/backend/agent/config/unified.py

@@ -39,12 +39,18 @@
 from ai.backend.common.config import BaseConfigSchema
 from ai.backend.common.configs import (
     EtcdConfig,
+    MemrayConfig,
     OTELConfig,
     PyroscopeConfig,
     ServiceDiscoveryConfig,
 )
 from ai.backend.common.configs.redis import RedisConfig
-from ai.backend.common.meta import BackendAIConfigMeta, CompositeType, ConfigExample
+from ai.backend.common.meta import (
+    NEXT_RELEASE_VERSION,
+    BackendAIConfigMeta,
+    CompositeType,
+    ConfigExample,
+)
 from ai.backend.common.typed_validators import (
     AutoDirectoryPath,
     GroupID,
@@ -2116,6 +2122,19 @@ class AgentGlobalConfig(BaseConfigSchema):
             composite=CompositeType.FIELD,
         ),
     ]
+    memray: Annotated[
+        MemrayConfig,
+        Field(default_factory=MemrayConfig),
+        BackendAIConfigMeta(
+            description=(
+                "Memray allocation-tracking configuration for the agent. "
+                "Enable it only while diagnosing agent memory growth: the capture file keeps "
+                "growing for as long as the process runs."
+            ),
+            added_version=NEXT_RELEASE_VERSION,
+            composite=CompositeType.FIELD,
+        ),
+    ]
     logging: Annotated[
         LoggingConfig,
         Field(default_factory=LoggingConfig),

src/ai/backend/agent/server.py

@@ -1,6 +1,7 @@
 from __future__ import annotations
 
 import asyncio
+import contextlib
 import functools
 import logging
 import os
@@ -1779,6 +1780,26 @@ def main(
         raise click.Abort() from e
 
     if not is_invoked_subcommand:
+        memray_tracker: contextlib.AbstractContextManager[Any] | None = None
+        if server_config.memray.enabled:
+            import memray
+
+            # `aiotools.start_server()` forks the service worker, so follow it to
+            # capture the allocations of the process that does the actual work.
+            # The PID suffix keeps a restart from colliding with an earlier capture,
+            # which memray would refuse to overwrite.
+            memray_destination = server_config.memray.output_destination
+            memray_destination = memray_destination.with_name(
+                f"{memray_destination.stem}.{os.getpid()}{memray_destination.suffix}"
+            )
+            memray_destination.parent.mkdir(parents=True, exist_ok=True)
+            memray_tracker = memray.Tracker(
+                memray_destination,
+                follow_fork=True,
+                native_traces=server_config.memray.native_traces,
+            )
+            memray_tracker.__enter__()
+
         server_config.agent_common.pid_file.write_text(str(os.getpid()))
         image_commit_path = server_config.agent_common.image_commit_path
         image_commit_path.mkdir(parents=True, exist_ok=True)
@@ -1823,6 +1844,8 @@ def main(
                 )
                 log.info("exit.")
         finally:
+            if memray_tracker is not None:
+                memray_tracker.__exit__(None, None, None)
             if server_config.agent_common.pid_file.is_file():
                 # check is_file() to prevent deleting /dev/null!
                 server_config.agent_common.pid_file.unlink()

src/ai/backend/common/configs/__init__.py

@@ -1,10 +1,12 @@
 from .etcd import EtcdConfig
+from .memray import MemrayConfig
 from .otel import OTELConfig
 from .pyroscope import PyroscopeConfig
 from .service_discovery import ServiceDiscoveryConfig, ServiceEndpointConfig
 
 __all__ = (
     "EtcdConfig",
+    "MemrayConfig",
     "OTELConfig",
     "PyroscopeConfig",
     "ServiceDiscoveryConfig",

src/ai/backend/common/configs/memray.py

@@ -0,0 +1,72 @@
+from __future__ import annotations
+
+from pathlib import Path
+from typing import Annotated
+
+from pydantic import AliasChoices, Field
+
+from ai.backend.common.config import BaseConfigSchema
+from ai.backend.common.meta import NEXT_RELEASE_VERSION, BackendAIConfigMeta, ConfigExample
+
+__all__ = ("MemrayConfig",)
+
+
+class MemrayConfig(BaseConfigSchema):
+    """Configuration for memray memory-allocation tracking.
+
+    Memray captures every allocation made by the process, which is useful for
+    diagnosing memory growth but adds significant runtime and disk overhead.
+    Enable it only while investigating a memory issue.
+    """
+
+    enabled: Annotated[
+        bool,
+        Field(default=False),
+        BackendAIConfigMeta(
+            description=(
+                "Whether to track memory allocations with memray. "
+                "The capture file grows continuously while the process runs, so keep this "
+                "disabled unless you are actively debugging memory usage."
+            ),
+            added_version=NEXT_RELEASE_VERSION,
+            example=ConfigExample(local="false", prod="false"),
+        ),
+    ]
+    output_destination: Annotated[
+        Path,
+        Field(
+            default=Path("./memray-output.bin"),
+            validation_alias=AliasChoices("output-destination", "output_destination"),
+            serialization_alias="output-destination",
+        ),
+        BackendAIConfigMeta(
+            description=(
+                "Path to store the memray allocation capture. "
+                "Memray refuses to start when the file already exists, so the process appends "
+                "its PID to the path. Ensure the parent directory exists and has enough space."
+            ),
+            added_version=NEXT_RELEASE_VERSION,
+            example=ConfigExample(
+                local="./memray-output.bin",
+                prod="/var/log/backend.ai/memray/agent.bin",
+            ),
+        ),
+    ]
+    native_traces: Annotated[
+        bool,
+        Field(
+            default=False,
+            validation_alias=AliasChoices("native-traces", "native_traces"),
+            serialization_alias="native-traces",
+        ),
+        BackendAIConfigMeta(
+            description=(
+                "Whether to also capture native (C/C++) call stacks, so that allocations made "
+                "inside extension modules are attributed to their native frames instead of the "
+                "calling Python frame. This adds runtime overhead and enlarges the capture, and "
+                "the native frames only resolve when debug symbols are available."
+            ),
+            added_version=NEXT_RELEASE_VERSION,
+            example=ConfigExample(local="false", prod="false"),
+        ),
+    ]