Repository navigation
perf: minimal AgentX 50k speed fixes (4 measured wins, +946/-18) - #6
Open
ajcasagrande wants to merge 6 commits into
Open
ajcasagrande wants to merge 6 commits into
ajcasagrande wants to merge 6 commits into
Conversation
A SIGUSR1 faulthandler dump of the live 50k run put the timing-manager thread inside warmup_credit_count -> warmup_credit_counts_by_lane -> the generator at trajectory_source.py:438, reached from _handle_warmup_return on *every* warmup credit return (agentic_replay.py:1584). Both properties rebuilt an O(lanes) tuple per call, so the warmup phase was quadratic: 53,740 returns x 50,000 lanes. Receipt: that single thread sat at 100% of a core with 8,500 voluntary context switches over 8 minutes (pure bytecode, not I/O), the pod used 1.17 of 192 cores, the Rust mock answered a direct request in 628 ms at 4 millicores, and throughput held at ~38-48 req/s with only ~8,000 of 50,000 connections ever opened. Nothing was blocked on work; the timing manager was just spinning on a redundant scan. Both values are fully determined by self.trajectories, whose per-lane snapshots are assigned only during _build_trajectories (the sole snapshot= site is the Trajectory construction at :656) and never reassigned, so caching on first use cannot go stale. (cherry picked from commit 9100f9739066d9472f636812c38a45a7d6b6ece2)
Tests and tooling assemble TrajectorySource with __new__ and set attributes directly, so the memo attributes added in the previous commit did not exist and the properties raised AttributeError. Detect the absent attribute with a sentinel and skip caching for those instances, so they keep recomputing and a later reassignment of trajectories cannot serve a stale value. (cherry picked from commit 76122f45eb1b148dd337be7f85f81096f83047e8)
SIGUSR1 faulthandler samples of the 50k timing_manager put 4 of 8 stacks in any_child_tracked_for_parent, reached from _finalize_failed_dispatches at the end of _spawn_children_and_register_gates on every child-spawning credit. It walked every entry of every tracked child, so the check was O(total tracked entries) per credit. Replace the plain dict with _ChildJoinMap, a dict subclass that keeps a parent -> entry-count index in step on every mutating path (setitem with overwrite, delitem, pop with and without default, popitem, setdefault, update, |=, clear) and drops parents at zero. The call site becomes an O(1) lookup. Because the index lives in the mapping, tests that assign into _child_to_join directly stay consistent with no edits. any_child_tracked_for_parent is kept as the brute-force test oracle; the new test checks the index equals it after deterministic and seeded random churn. Signed-off-by: Anthony Casagrande <acasagrande@nvidia.com> (cherry picked from commit 8fc127fd0dcdadefc91c7ce1de570c1b872c3379)
… channel Every streamed chunk crossed that channel through pydantic (model_dump(mode="json")) + orjson.dumps on the worker, then orjson.loads + a four-member union validation in the record processor. With ~200 chunks per request at ~9.6k req/s that is ~1.9M chunks/s paying both ends. Replace only the wire encoding with typed msgpack, mirroring the existing ZMQStreaming push/pull clients: a pre-built msgspec Encoder with msgspec_enc_hook on the send side and a typed Decoder on the receive side. SSEMessage/SSEField are slots dataclasses, so msgspec handles them natively in C and the decoded objects are the same types the JSON path produced. Non-SSE responses (text, binary, AWS eventstream, mixed lists) travel as tagged structs in their original order. The per-record remainder is still dumped with mode="json" and re-validated, so its semantics are unchanged. The codec belongs to the channel, not the caller: WIRE_CODECS maps (client type, address) to codec kwargs and BaseZMQCommunication.create_client injects them, so one socket can never carry two encodings and no flag can quietly put JSON back on it. A cross-encoding frame fails loudly in both directions and both are tested. RECORDS and every other channel stay on JSON. Estimated (not yet measured) saving is 3-6 us per chunk, ~6-12 cores pod-wide. Port note: master's `create_client` caches on (type, address, bind) with no kwargs-aware key or `attach_lifecycle`, so the codec injection is placed ahead of master's cache check and master's caching is otherwise unchanged. Master has no `AwsEventStreamMessage` response type and no `msgspec_enc_hook` (both came from upstream changes outside this port), so the AWS wire variant and its test case are dropped and the encoder is a plain `msgspec.msgpack.Encoder()` like the other msgpack clients here; the frame only carries builtins and slots dataclasses, so no hook is reached. Signed-off-by: Anthony Casagrande <acasagrande@nvidia.com> (cherry picked from commit 28c8c4ebd072f8a15dd7d27ec2c6fb2e05338f5b)
…atus replies A SIGUSR1 sample of the 50k dataset_manager sat in router_reply_client._wait_for_response -> Message.to_json_bytes -> pydantic model_dump, with GC visible. The dataset ROUTER channel carries only three request types, and after a worker opens its local mmap client it never sends ConversationRequest again (the conversation cache it needs only exists once the client is open). The one reply that is both repeated and large is DatasetConfigStatusResponse: every not-yet-configured worker polls it each STATE_POLL_INTERVAL, and once configured each reply carries the full DatasetMetadata, which was re-dumped through pydantic per poll. Encode the notification once per published notification object in the DatasetManager and splice those bytes into the reply envelope. notification is the response's last field, so the spliced frame is byte-identical to the plain to_json_bytes() output; the test asserts that, one encode across repeated polls, and re-encoding when a new notification object is published. No wire change: receivers decode exactly as before. Signed-off-by: Anthony Casagrande <acasagrande@nvidia.com> (cherry picked from commit 4931e103eb09f83ba7c95d0498e4bf564013303f)
Signed-off-by: Anthony Casagrande <acasagrande@nvidia.com> (cherry picked from commit 048f38a0b0208da7ec68963e788b5ef2b8702810)
Try out this PRQuick install: pip install --upgrade --force-reinstall git+https://github.com/ai-dynamo/aiperf.git@3e555153cf7a55b9ce7d296753e82443eabbf8eeRecommended with virtual environment (using uv): uv venv --python 3.12 && source .venv/bin/activate
uv pip install --upgrade --force-reinstall git+https://github.com/ai-dynamo/aiperf.git@3e555153cf7a55b9ce7d296753e82443eabbf8eeLast updated for commit: |
This branch has not been deployed
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.
Minimal set of speed fixes for AgentX at 50,000 concurrency, on agentx-harness
master(60a094f). Every commit here has an individually measured number; unmeasured and inert changes were deliberately excluded.Measured at 50k, official
semianalysis_cc_traces_weka_062126corpus (393 traces), 100 workers, 32 record processors.Commits
7a89e7d8dbe839d1c8__new__-built sources raiseAttributeError98638157d974a52da6b4b0bed643e555153cWhy the first one matters most
warmup_credit_countrebuilt an O(lanes) tuple on every credit return, so warmup was quadratic: 53,740 returns x 50,000 lanes. At 50k this pinned the single timing-manager thread at a full core while the pod used 1.17 of 192 cores, holding throughput to ~38 req/s.Excluded on purpose
payload_fragments(+175 lines) — its cache is disabled by default, so it would ship inert.Wire bytes
Client -> server HTTP bodies are unchanged. The one format change is AIPerf's internal ZMQ channel between worker and record-processor (
974a52da6), which is not the benchmark traffic. If you want the smallest possible diff, dropping that commit leaves ~720 lines across 4 files carrying the 12x plus the two smaller wins.Not verified
Not benchmarked as a branch, and not independently reviewed — the numbers above are per-commit measurements, not a measurement of this stack. All commits GPG-signed.