feat: add opt-in memray allocation tracking (agent, manager, storage-proxy, app-proxy)#12800
Open
Yaminyam wants to merge 8 commits into
Open
feat: add opt-in memray allocation tracking (agent, manager, storage-proxy, app-proxy)#12800Yaminyam wants to merge 8 commits into
Yaminyam wants to merge 8 commits into
Conversation
Add a `[memray]` config section to the agent so its memory allocations can be captured for diagnosing memory growth, mirroring the existing app-proxy coordinator profiling setup. - New shared `MemrayConfig` (common/configs/memray.py) with `enabled`, `output-destination`, and `native-traces`, so manager/storage can reuse it. - The agent installs a `memray.Tracker(follow_fork=True, native_traces=...)` in `main()` before the service worker forks, and tears it down in the `finally` block. `follow_fork=True` captures the forked worker that does the real work; the capture path gets the PID appended so a restart does not collide with an earlier capture (memray refuses to overwrite). - Disabled by default; enabling adds runtime/disk overhead. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
github-actions
Bot
added
size:L
Yaminyam
marked this pull request as draft
July 13, 2026 08:08
Extend the shared MemrayConfig to the manager, mirroring the agent. - Add `memray` to BootstrapConfig (read in main() before the workers fork) and to ManagerUnifiedConfig (for the generated sample). - Install `memray.Tracker(follow_fork=True, native_traces=...)` in main() before `aiotools.start_server()` so every forked worker is captured; tear it down in the finally block. PID is appended to the capture path so a restart does not collide with an earlier capture. - Make the shared prod example path component-neutral (output.bin). Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
Migrate the app-proxy coordinator and worker to the shared MemrayConfig
via a new `[profiling.memray]` table, non-breaking.
- Add nested `memray: MemrayConfig` to ProfilingConfig; consumers read
`profiling.memray.{enabled,output_destination,native_traces}` (adds
native-traces support to app-proxy).
- Keep `enable_memray` / `memray_output_destination` working as deprecated
aliases; a model_validator folds them into `memray` when the new table is
absent (an explicit `[profiling.memray]` wins).
- Regenerate the coordinator/worker samples. The coordinator sample also
picks up previously un-regenerated schema fields (pool_recycle,
pool_pre_ping, traefik client_ip_strategy) that were already merged.
Note: app-proxy uses snake_case TOML keys, so the shared memray table renders
its kebab serialization aliases (output-destination, native-traces); both
snake and kebab are accepted on load.
Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
Yaminyam
changed the title
Member
|
From what I remember, the memray option added to the App Proxy Coordinator was intended to be temporary and is expected to be removed. I think it's worth confirming whether that's still the plan. @HyeockJinKim |
The worker's capture -- the only one worth reading, since the workers do the real work -- was never closed, so it ended up truncated and unparseable: memray stopped at the first bad record, reporting a 0.0s span out of a 366 MB file. Two things line up to cause it: - memray's `Tracker` writes the capture's trailer only from `__exit__`. It has no `__del__` and registers no `atexit` handler of its own. - We only called `__exit__` in the master's `finally`, which never runs in a forked child. An `atexit` handler in the child does not help either: `aiotools` ends its forked children with a bare `os._exit()` (`aiotools/fork.py`), which runs neither `atexit` handlers nor interpreter finalization. So the worker has to close the tracker itself, explicitly, while it is still running Python. Hoist the tracker lifecycle into two helpers next to the Pyroscope profiler -- `start_memray_tracker()` for the master, `close_memray_tracker_in_worker()` for the child -- and call the latter from the shutdown path of each component's forked entry point (agent, manager, app-proxy coordinator and worker, all of which fork via `aiotools.start_server()`). `Tracker.__exit__` is not idempotent (a second call raises), so the master and each worker close exactly one tracker: their own. The module global is a per-process copy after the fork, so the worker clearing it does not disturb the master. This also gives app-proxy the PID-suffixed capture path the agent and manager already had, so a restart no longer collides with an earlier capture (memray refuses to overwrite one). Verified against a live agent: the forked worker's capture goes from a 0.0s span that memray rejects with "Invalid record type" to a finalized 68.7s capture that parses cleanly. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
Extend the shared MemrayConfig to the storage proxy, matching the agent, manager, and app-proxy. The storage proxy also forks its service worker via `aiotools.start_server()`, so the same follow-fork capture applies. - Add `memray` to StorageProxyUnifiedConfig (read in main() before the worker forks) and regenerate the sample. - Install `memray.Tracker(follow_fork=True, ...)` in main() via the shared helper; the master closes it in the finally block and each forked worker closes its own capture from its shutdown path (aiotools ends children with a bare os._exit(), so atexit never runs). Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
github-actions
Bot
added
the
comp:storage-proxy
…ngelog Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
Yaminyam
changed the title
Yaminyam
marked this pull request as ready for review
July 15, 2026 04:39
|
Oh cool, this adds some really neat performance insights! 🚀😌 Maintainer? Turn off weaves from non-maintainers → |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
Adds an opt-in memray allocation-tracking option across the agent, manager, storage proxy, and app-proxy, backed by a shared
MemrayConfig(common/configs/memray.py). Useful for diagnosing memory growth. Mirrors the profiling setup the app-proxy already shipped (memray.Tracker(follow_fork=True)).Shared config
MemrayConfigfields:enabled(defaultfalse)output-destination(default./memray-output.bin)native-traces(defaultfalse) — also capture native C/C++ stacks (memray --native); higher overhead, needs debug symbols to resolvePer-component
[memray]section onAgentGlobalConfig; installsmemray.Tracker(follow_fork=True, native_traces=...)inmain()before the service worker forks.follow_fork=Trueis required becauseaiotools.start_server()forks the worker that does the real work. The capture path gets the PID appended (memray refuses to overwrite an existing file).memrayonBootstrapConfig(read inmain()before the workers fork) and onManagerUnifiedConfig(for the generated sample).[memray]onStorageProxyUnifiedConfig, installed inmain()before the worker forks.[profiling.memray]table, non-breaking: the oldenable_memray/memray_output_destinationkeys keep working as deprecated aliases, reconciled by amodel_validatorwhen the new table is absent (an explicit[profiling.memray]wins). This also adds native-traces support and the PID-suffixed capture path to app-proxy.Disabled by default everywhere; enabling adds runtime and disk overhead (the capture grows for the process lifetime).
Worker-side capture close (fix)
The forked worker's capture — the only one worth reading, since workers do the real work — used to end up truncated and unparseable (memray reported a 0.0s span out of a 366 MB file):
Trackerwrites the capture trailer only from__exit__; it has no__del__and registers noatexithandler.finally, which never runs in a forked child — andaiotoolsends children with a bareos._exit(), soatexitwould not run either.The tracker lifecycle now lives in two helpers next to the Pyroscope profiler:
start_memray_tracker()for the master andclose_memray_tracker_in_worker()called from each component's forked shutdown path.Tracker.__exit__is not idempotent, so master and worker each close exactly one tracker: their own (the module global is a per-process copy after the fork).Notes
output-destination,native-traces); both snake and kebab are accepted on load.pool_recycle,pool_pre_ping, traefikclient_ip_strategy) that were already merged — the sample is a generated artifact, so it is committed in sync.Verification
/metrics-> 200,follow_forkcapturing the worker + kernel-runner forks (hundreds of thousands of allocations),native-traces=truepath confirmed via capture metadata. Rolled back after testing.pants fmt/lint/checkgreen across all changed files; component samples regenerated from schema.🤖 Generated with Claude Code