Speed up serialization and remove the quadratic metadata scans - #74
Merged
Merged
Conversation
The protobuf Python API converts every element of a repeated double field to
and from a boxed Python float. For bulk arrays that dominates: moving 200k
doubles cost 27.4 ms to encode and 7.9 ms to decode, per direction, against
about 2 ms for the gRPC round trip carrying them.
None of that work is required by the format. A packed repeated double is
encoded on the wire as a tag, a length, and a contiguous buffer of
little-endian doubles, which is byte for byte what NumPy already holds. So
this reads and writes that buffer directly:
- set_array_data builds the Array with its metadata, then merges in a
hand-framed field-6 blob, which protobuf stores contiguously
- get_array_data re-serializes the message and reads the packed payload
with np.frombuffer, which is two memory copies rather than N conversions
Output is byte-identical to what the protobuf API produces, verified across
sizes and metadata combinations and in both directions against messages built
the ordinary way. The protocol does not change and peers in any language are
unaffected: verified against a server built from 510d3f8, in both directions,
for a 5000-element array, a scalar, and partials.
Encode wins at every size, including a single element, so it is
unconditional. Decode has a fixed cost of about 1.9 us from the re-serialize,
so below 64 elements it stays on the protobuf container. The bound is already
in hand from the message's start and end indices, so the check is nearly free
and scalar-heavy disciplines are unaffected.
Coercion the protobuf API used to do for us is now explicit: values are cast
through little-endian float64, which pins byte order on a big-endian host and
stops a float32 array being written at half length or an integer array having
its bit pattern reinterpreted as doubles.
Measured, per direction for 200k doubles: encode 27.4 -> 0.15 ms, decode
7.9 -> 0.11 ms. End to end over the streaming transport, a 250k-element round
trip goes from 121 ms to 36 ms; the remainder is the default 1000-double
chunking, which this does not touch.
DisciplineServer.preallocate_partials and DisciplineClient._recover_partials
both looked up a variable's shape by scanning the entire variable metadata
list, building a throwaway list of every match and then indexing [0]. They did
it twice per declared partial, once for the function and once for the
variable, so the work was quadratic in the size of the discipline. It was paid
on every gradient call, not once at setup.
Both now build a {name: shape} index once per call and look up through it.
Measured with 100 variables and 100 declared partials: preallocate_partials
drops from 4.11 ms to 0.157 ms, and a full gradient round trip from 23.0 ms to
14.6 ms. The remaining time is the per-message gRPC cost, which is a function
of the message count rather than of anything computed here.
The two sites also carried separate copies of the rule for the shape of a
Jacobian block, including the edge cases where the function or the variable is
scalar. Since both copies had to be rewritten anyway, the rule now lives in
philote_mdo.utils.get_partials_shape, so the two cannot drift.
get_chunk_indices derived its answer by building an array of chunk starts with np.arange, then an array of ends with np.append, then zipping the two. For a variable that fits in a single chunk, which is every scalar and most small arrays, that is two array constructions to produce one pair. It now returns that pair directly. The guard is 0 < num_values <= chunk_size, so an empty variable still yields no chunks rather than one empty chunk, which is what the array path produced. Verified identical to the previous derivation across every combination of num_values in 0..59 and chunk_size in 1..11. Across 100 scalar variables, 0.095 ms to 0.005 ms. Small on its own, but it is on the path of every variable of every call.
With the per-element conversion cost gone from encoding, the runtime of a stream is set by how many messages it carries, not by how large they are. A bare gRPC echo servicer with no discipline logic costs about 64 microseconds per message per direction, against roughly 5 microseconds of Philote's own work to build and read one. So the default wants to be as large as the message ceiling safely allows, and 1000 doubles was far below it. For a 250k-element round trip, measured across the range: num_double=1000 250 msgs/dir 41.7 ms num_double=10000 25 msgs/dir 6.1 ms num_double=100000 3 msgs/dir 2.9 ms num_double=250000 1 msg /dir 3.0 ms The curve is flat past 100k, so the ceiling is what picks the value rather than the measurement. A 100k-double chunk serializes to about 780 KiB, a fifth of gRPC's 4 MiB default limit, which leaves room both for message metadata and for a peer that has lowered the limit. 500k would sit at 95% of it. This is a default only. StreamOptions is negotiated over SetStreamOptions exactly as before, both sides move together, and anyone who sets num_double explicitly is unaffected. Variables smaller than a chunk are untouched, since chunking only splits what exceeds the size. The unit test that pinned the old default has been updated.
add_input, add_output, add_discrete_input and add_discrete_output each
rejected a duplicate declaration by scanning the entire metadata list for a
matching name and type. Declaring a discipline was therefore quadratic in its
own size, and every variable paid the cost of all the ones before it.
The declared (type, name) pairs now live in a set alongside the lists, so the
check is a hash lookup. _clear_data resets it with the rest of the metadata.
variables before after
100 2.0 ms 0.3 ms
500 43.6 ms 1.7 ms
1000 171.6 ms 3.5 ms
2000 684.4 ms 6.7 ms
Setup runs once per client rather than once per call, so this does not show up
in a compute benchmark. It shows up when a large discipline starts, and for an
OpenMDAO group of remote components it is paid per component.
The rule the scan implemented is unchanged: the check is per variable type, so
an input and an output may still share a name, and the residual that an
implicit output implies still does not collide with it. Both of those, and
redeclaration after _clear_data, now have tests; the four duplicate-rejection
cases already did.
SetVariableShapes on the server and send_variable_shapes on the client both
resolved each incoming shape by scanning the whole variable metadata list, and
scanned it a second time to find the residual entry that an implicit output
implies. With one message per dynamic variable, that is quadratic in the size
of the discipline.
Both now build a {(type, name): variable} index once before consuming the
messages. setdefault keeps the first entry for a key, matching the break in
the loop it replaces.
variables before after
100 20.4 ms 12.5 ms
250 82.0 ms 30.9 ms
500 264.1 ms 61.5 ms
1000 922.4 ms 119.9 ms
What remains is the per-message gRPC cost of the stream itself, which is now
linear in the number of dynamic variables.
The server's two rejections are unchanged and still ordered the same way: an
unknown variable is reported as not found, and a known one that was not
declared dynamic is rejected as not allowing dynamic shapes. Both already had
tests, as does the residual propagation.
npm audit reported 42 findings against the Docusaurus toolchain, 2 critical
and 26 high. This takes it to 17.
- npm audit fix updated the webpack-dev-server, sockjs, ws and
websocket-driver chains within the ranges their parents already allow,
which cleared both criticals
- serialize-javascript and uuid had no fix path through their parents, so
they are pinned through overrides to ^7.1.0 and ^11.1.1
The remaining 17 are all image-size, reached through @docusaurus/mdx-loader.
Its advisory range is *, meaning every published version including the latest
is affected, so there is nothing to upgrade to. It parses image dimensions at
build time, and the only images in this repository are a checked-in logo and
favicon.
All of these are build-time dependencies. The deployed site is static HTML, so
the exposure is a developer running the dev server and the CI build, not
anyone reading the documentation.
Docusaurus itself stays at 3.9.2. 3.10.2 was tried first and rejected: it left
the audit count unchanged at 24, and it breaks the build, because the site
sets future.v4 and 3.10 requires the @docusaurus/faster package for it.
Verified that npm ci reproduces the tree from the lockfile and that npm run
build succeeds from that clean install, since the docs deploy and release
workflows both depend on it.
The encoding entry quoted a round-trip figure that the later chunk size change supersedes, which read as a contradiction against the figure in that entry. It now says which chunk size it was measured at and points forward. Bullets within a section are contiguous elsewhere in the file; these were separated by blank lines.
Codecov Report✅ All modified and coverable lines are covered by tests. 📢 Thoughts on this report? Let us know! |
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
1 participant
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.
Performance work on the existing API. No
.protochanges, no wire format changes, no submodule move — verified by running the new code against a server and client built from510d3f8, in both directions.Results
Commits
Each is independently revertable.
17b6b0e— read and writeArray.datathrough the packed wire buffer instead of protobuf'srepeated doubleAPI, which boxes every element as a Python float. A packedrepeated doubleis a tag, a length, and a contiguous block of little-endian doubles, which is byte for byte what NumPy already holds, so the emitted bytes are identical. Decoding stays on the protobuf container below 64 elements, where re-serializing the message to reach its payload costs more than it saves.d3e8ebe—preallocate_partialsand_recover_partialsrescanned the whole variable metadata list twice per declared partial, on every gradient call. Now indexed by name. The duplicated Jacobian block shape rule moved toutils.get_partials_shape().7a750f7—get_chunk_indicesreturns the single-chunk case directly rather than deriving it through two NumPy array constructions.4357357— defaultnum_double1,000 → 100,000. With the encoding cost gone, a stream's runtime is set by its message count, at ~64 µs per message per direction. A chunk is ~780 KiB against gRPC's 4 MiB ceiling. Default only;StreamOptionsis negotiated as before.7a27e5d— the fouradd_*methods checked for duplicates by scanning the metadata list, making discipline declaration quadratic in its own size. Now a set of declared(type, name)pairs.9b5b854—SetVariableShapesandsend_variable_shapesresolved each shape by scanning the list, twice for an implicit output. Now indexed by type and name.8efcb16— docs site npm advisories: 42 (2 critical, 26 high) → 17. See the caveat below.6392a2f— changelog consistency.Where the floor is now
Everything remaining is gRPC's per-message cost. A bare echo servicer with no discipline logic costs 32.3 ms for 250 messages; Philote's own work is ~5 µs per message. So for bulk arrays you cut message count via chunk size, but a discipline with many small variables is bounded at ~64 µs per variable per direction. Reducing that means packing multiple variables into one message, which needs a protocol change and is deliberately out of scope here.
Caveats
image-sizevia@docusaurus/mdx-loader. Its advisory range is*— every published version is affected, so there is nothing to upgrade to. All are build-time; the deployed site is static. Docusaurus stays at 3.9.2: 3.10.2 left the audit count unchanged and breaks the build, since the site setsfuture.v4and 3.10 requires the@docusaurus/fasterpackage.test_discipline_client.py::test_init, which pinned the oldnum_doubledefault.for jac, value in jac.items(), rebinding the loop variable over the dict being iterated. It works only because the items view is bound before the first rebind. Pre-existing, out of scope, tracked upstream as Loop variable shadows the dict being iterated in both gradient RPCs #71.Testing
309 tests pass. New coverage: 23 tests for the encoding helpers asserting byte identity against messages built the ordinary way; 4 for
get_partials_shape; chunk boundary and empty cases forget_chunk_indices; and two behaviours the duplicate-check rewrite could have broken and nothing tested — that an input and an output may share a name, and that_clear_datapermits redeclaration.Beyond the suite: cross-version interop against a
510d3f8worktree in both directions, exact values for a 5000-element array, a scalar, and partials; a differential check ofget_chunk_indicesagainst its previous derivation across all 660 combinations ofnum_values0–59 ×chunk_size1–11; andnpm cifollowed bynpm run buildfrom a clean install, since the docs deploy and release workflows depend on both.