[SYSTEMDS-3946] Enable sending of large (>2GiB) FederatedRequests and…

[Biranavan-Parameswaran]

Federated transfers previously failed for payloads above 2 GiB. Netty's frame length field is a signed 32-bit integer, so a single request or response is capped at Integer.MAX_VALUE bytes.

This patch adds a streaming chunked codec. The sender splits a serialized payload into bounded frames and the receiver reassembles them, so the size of one logical message on the wire is no longer limited by a single frame.

Wire format

Every message is preceded by a one byte format marker, so the receiver knows which decoder to use: MARKER_OBJECT_ENCODER or MARKER_CHUNKED. A chunked message then follows as a sequence of frames:

1 byte   frame type (TYPE_DATA, TYPE_END, TYPE_ERROR)
4 bytes  payload length, big endian
n bytes  payload

FederatedChunkEncoder serializes on a pool thread and feeds frames through Netty's ChunkedWriteHandler. FederatedChunkDecoder reassembles frames into an ObjectInputStream that the deserializer consumes as they arrive. Neither side ever materializes the whole payload as a single array, while transferring.

Backpressure runs in both directions: a sender that outruns the socket and a receiver that outruns the deserializer each pause until the other side catches up, so a 4 GB message still holds only about 64 MB in memory at a time when transferring.

Routing

FederatedFormatEncoder.useObjectEncoder() uses the legacy ObjectEncoder only when the lineage cache is active and the message is a lineage cacheable FederatedResponse below STREAM_THRESHOLD, which is 2000 MiB or 1.953 GiB. The
lineage cache is off by default, so in a default configuration every message takes the chunked path.

Why the legacy encoder is kept at all: the lineage serialization cache (LineageCache.putSerializedObject) stores an entire serialized response as one INT_MAX bounded byte[] and only the ObjectEncoder path produces that array. Streaming never materializes it. So the legacy path survives for exactly the case that requires it and for nothing else.

A size guard is still needed because that byte[] cannot exceed INT_MAX, so a response too large for it has to stream whether or not it is cacheable.

The guard is deliberately 2000 MiB rather than INT_MAX - 1. It sees an estimate of the raw payload, but the wire adds a measured 0.4883% of framing, so a payload gated at INT_MAX - 1 overflows by about 10 MB. The wire was measured to reach Integer.MAX_VALUE at a raw payload of about 1.990 GiB, and 2000 MiB sits around 38 MiB under that.

Experiments

One client and one worker on the same host, timed as the client's Total elapsed time over an elementwise operation on a federated matrix, so the matrix crosses the codec. AMD Ryzen 5 4600H, 12 cores, 11 GB RAM, OpenJDK
17.0.10 on WSL2, both JVMs at -Xmx8g -Xms1g -Xmn256m.

payload	object encoder	chunked	ratio
80 KB	1.398 ± 0.059 s	1.439 ± 0.065 s	0.97x
8 MB	2.218 ± 0.080 s	2.238 ± 0.074 s	0.99x
512 MB	86.142 ± 3.802 s	20.527 ± 2.138 s	4.20x faster
1 GiB	324.350 ± 6.497 s	64.478 ± 5.273 s	5.03x faster

n = 15 for 80 KB and 8 MB, n = 3 for 512 MB and 1 GiB. Small payloads require extra runs since their effect size matches the background noise between runs. At 512 MB and 1 GiB the gap is 4x to 5x, far outside the spread, and one 1 GiB pair already costs about 6.5 minutes.

Small payloads: no measurable regression. Chunked is 2.9 % slower at 80 KB and 0.9 % at 8 MB, both within the run to run noise. The overhead is fixed per message, so it does not grow with the payload.

Large payloads. The object encoder holds the whole message as one contiguous buffer at both ends and cannot start sending before serialization finishes. The chunked codec overlaps the two and never holds more than the frame queue, so cost stays close to linear and memory stays bounded.

Different baseline. These figures supersede an earlier table in this thread, which compared against an outdated baseline.

Chunk size

4 MB (1 << 22), from a sweep over 256K, 1M, 4M and 8M across four payloads, 48 runs. Median fed_+ in seconds, the federated instruction time, not the full round trip of the table above:

Median fed_+ against chunk size, as a percentage faster than the 256 KB chunk, one line per payload. 80 KB is not drawn, it is one frame at every chunk size.

payload	frames at 256K / 1M / 4M / 8M	256K	1M	4M	8M
0.08 MB	1 / 1 / 1 / 1	0.755	0.787	0.900	0.835
8 MB	31 / 8 / 2 / 1	1.208	1.165	1.099	1.066
512 MB	1954 / 489 / 123 / 62	10.707	10.203	10.013	10.443
1074 MB	4096 / 1024 / 256 / 128	17.439	16.020	15.578	15.917

4 MB wins at 512 MB and at 1074 MB and holds as the payload doubles, while 8M gives part of that back (4.3 % slower at 512 MB, 2.2 % at 1074 MB) and 256K pays too much framing overhead. 8M leads the 8 MB row only because the payload fits in one frame there, and 0.08 MB is one frame at every size, so neither row says anything about chunking. 4 MB also caps buffering in flight at 64 MB.

Above the 2 GiB cap

The range the codec exists for. One run per size, on a different macOS host with about 20 GB of usable RAM and a lighter workload, so these absolute numbers do not compare with the round trip table further up. The trend within the
table is the point:

Peak resident memory of the worker against payload size, above the 2 GiB single frame cap with a 2x payload reference line. Sampled every 0.4 s, 1 run per payload.

payload	round trip	throughput	peak worker RSS
2.2 GB	11.56 s	190 MB/s	5961 MB
2.6 GB	14.26 s	182 MB/s	7689 MB
3.0 GB	16.14 s	186 MB/s	6744 MB
4.0 GB	25.74 s	155 MB/s	9795 MB

All correct, 0 exceptions, 0 OOM. Throughput stays flat and peak worker RSS stays at 2.2x to 3.0x the payload, which is the workload's own two copies, the received matrix plus the result, with no serialization buffer on top. The object encoder
cannot reach this range at all, it fails at about 1.99 GiB.

Known limitations

• Every measurement above is a sequential single client round trip. Because routing now sends essentially all traffic through the chunked path with two extra pool tasks per message, behavior under many concurrent small requests is not yet measured.
• The lineage serialization cache still requires the legacy encoder, so the ObjectEncoder path cannot be removed until that cache can accept a streamed response.

[ywcb00]

Thank you very much for the PR @Biranavan-Parameswaran :)
I left some minor comments in the code. Could you please have a look at it and resolve it if you find the time. Thanks.

[ywcb00]

We should use the regular encoder as long as we can, i.e., up to the largest possible message size. Can we increase this default threshold from 1.5GB to (INT_MAX - 1) bytes?

[Biranavan-Parameswaran]

Good suggestion, but INT_MAX-1 is actually unsafe here. Routing uses a size estimate, not the exact wire size, and the ObjectEncoder overflows its Integer.MAX_VALUE ByteBuf once serialization framing is added. I measured that cliff at about 1.990 GiB, so an estimate just under INT_MAX can still overflow on the wire. I set the threshold to 2000L << 20 (about 1.953 GiB), roughly 40 MB under the cliff, so the regular encoder is used as high as we safely can.

Note the routing also changed since your review. Responses now route by lineage cacheability instead of size, so STREAM_THRESHOLD is now the size guard on the object encoder path rather than the main router.

[Biranavan-Parameswaran]

Thanks for the review @ywcb00. All three comments are addressed, and I also ran experiments to back the design choices. Summary of what changed since the reviewed commit.

Comments

• Threshold (1): INT_MAX-1 can overflow on the wire because routing uses a size estimate. I measured the overflow cliff at about 1.990 GiB and set the threshold about 40 MB under it (2000L<<20). Full reasoning in the inline reply.
• Redundant catch (2): removed.
• Rename (3): FederatedFormatDetector is now FederatedFormatDecoder.

Follow up work

• Chunk size 4 MB (1<<22). Swept 256K, 1M, 4M and 8M across four payloads, 48 runs. Median round trip in seconds:

  payload	256K	1M	4M	8M
  0.08 MB	0.755	0.787	0.900	0.835
  8 MB	1.208	1.165	1.099	1.066
  512 MB	10.707	10.203	10.013	10.443
  1074 MB	17.439	16.020	15.578	15.917

  4 MB is the robust optimum. It wins on large and xl and stays stable as the payload doubles. 8M starts to regress on big payloads and 256K pays too much framing overhead. Small payloads are one frame so they are chunk size indifferent. It also caps in flight buffering at 16 x 4 MB, so 64 MB.

• Streaming vs object encoder. Mean elapsed in seconds, object encoder vs chunked:

  payload	object encoder	chunked
  tiny	1.36	1.46
  mid	2.16	2.21
  512 MB	74.1	18.5
  1 GiB	305.5	29.9

  Chunked matches the object encoder on small payloads and is 4x to 10x faster on large ones.

• Scaling past 2 GB. Swept 2.2 to 4.0 GB, all correct with 0 exceptions and 0 OOM:

  payload	round trip s	MB/s	peak worker MB
  2.2 GB	11.56	190	5961
  2.6 GB	14.26	182	7689
  3.0 GB	16.14	186	6744
  4.0 GB	25.74	155	9795

  Throughput stays flat around 155 to 190 MB/s and worker peak memory tracks payload at about 2.2 to 3.0x with no second full payload buffer. The old object encoder crashes near 1.99 GiB, so this is the range the codec exists for.

• Routing now based on lineage cacheability instead of size. Lineage cacheable responses use the object encoder, everything else streams. New FederatedFormatRoutingTest covers the four cases.

• Close race fix. The worker CloseListener treated a client close during a chunked response drain as a write failure and tried a second write on the dead channel. It now treats an inactive channel mid stream as a normal end of stream.

• Slimmed FederatedChunkEncoder to a factory and tightened the routing predicate.

[codecov]

Codecov Report

❌ Patch coverage is 80.82192% with 42 lines in your changes missing coverage. Please review.
✅ Project coverage is 71.59%. Comparing base (fa9824c) to head (60c5c24).
⚠️ Report is 3 commits behind head on main.

Files with missing lines	Patch %	Lines
...ontrolprogram/federated/FederatedChunkEncoder.java	79.74%	14 Missing and 2 partials ⚠️
...ontrolprogram/federated/FederatedChunkDecoder.java	81.57%	10 Missing and 4 partials ⚠️
...ntrolprogram/federated/FederatedWorkerHandler.java	54.54%	4 Missing and 1 partial ⚠️
...ntrolprogram/federated/FederatedFormatEncoder.java	85.71%	2 Missing and 1 partial ⚠️
...ntrolprogram/federated/FederatedFormatDecoder.java	81.81%	1 Missing and 1 partial ⚠️
...untime/controlprogram/federated/FederatedData.java	92.85%	0 Missing and 1 partial ⚠️
...me/controlprogram/federated/FederatedResponse.java	0.00%	0 Missing and 1 partial ⚠️

Additional details and impacted files

@@             Coverage Diff              @@
##               main    #2496      +/-   ##
============================================
- Coverage     71.60%   71.59%   -0.01%     
- Complexity    50259    50390     +131     
============================================
  Files          1623     1631       +8     
  Lines        194314   194856     +542     
  Branches      37965    38025      +60     
============================================
+ Hits         139130   139506     +376     
- Misses        44277    44413     +136     
- Partials      10907    10937      +30     

☔ View full report in Codecov by Harness.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.
• 📦 JS Bundle Analysis: Save yourself from yourself by tracking and limiting bundle sizes in JS merges.

[ywcb00]

Thank you very much for this contribution. :)
The code looks very good. I only added some minor comments regarding variable naming.
A general aspect to consider: Do you think it is possible to make the FederatedChunkCodecTest even faster in terms of runtime needed for the test while testing the same functionality?
All the best,
David

[ywcb00]

How did you decide on this number here? Was it chosen arbitrarily or are there certain factors/reasons behind this decision?

[Biranavan-Parameswaran]

It's randomly chosen because the queue never fills up locally to test it.

Happy to adjust or make it configurable.

[Biranavan-Parameswaran]

Thank you very much for this contribution. :) The code looks very good. I only added some minor comments regarding variable naming. A general aspect to consider: Do you think it is possible to make the FederatedChunkCodecTest even faster in terms of runtime needed for the test while testing the same functionality? All the best, David

Thanks @ywcb00 , all comments resolved! :)

Regarding FederatedChunkCodecTest: I tried a few variants. The actual work is only ~50ms. I have a slightly faster version, but it costs a lot of readability in the polling helpers. Happy to commit it if you'd prefer.