Skip to content

[CELEBORN-2442] Optimize meta memory for very large partitions - #3825

Open
FMX wants to merge 2 commits into
apache:mainfrom
FMX:optimize-meta-memory
Open

[CELEBORN-2442] Optimize meta memory for very large partitions#3825
FMX wants to merge 2 commits into
apache:mainfrom
FMX:optimize-meta-memory

Conversation

@FMX

@FMX FMX commented Aug 24, 2026

Copy link
Copy Markdown
Contributor

What changes were proposed in this pull request?

This PR reduces the memory footprint of PartitionLocation metadata while preserving the existing wire protocol and public getter behavior.

The main changes are:

  • Introduce an immutable, weakly interned WorkerEndpoint to share worker host and port information across partition locations.
  • Move the cached host:pushPort and host:fetchPort strings into the shared endpoint.
  • Lazily initialize StorageInfo and RoaringBitmap using thread-safe initialization.
  • Keep getMapIdBitMap() non-null for compatibility, and add a non-materializing accessor for serialization, HTTP API, and read-side inspection.
  • Avoid allocating empty bitmaps when locations are created by LifecycleManager and SlotsAllocator.
  • Preserve the current storageTypes.head selection behavior in LifecycleManager.
  • Keep string keys in WorkerPartitionLocationInfo to avoid parsing and boxing overhead on the Worker hot path.
  • Add allocation-free helpers for converting partition location IDs into packed long values.
  • Avoid materializing optional metadata from logging, protobuf serialization, HTTP rendering, and range-read paths.
  • Re-intern WorkerEndpoint instances after same-version Java deserialization.
  • Fix peer bitmap serialization so that the peer protobuf contains the peer's bitmap instead of the primary location's bitmap.
  • Make range reads fall back to the peer bitmap and fail open when bitmap metadata is unavailable.
  • Add JOL-based manual memory benchmarks and focused unit tests.

No protobuf schema is changed by this PR.

Why are the changes needed?

A shuffle may contain millions of PartitionLocation instances. Previously, every instance independently retained:

  • Worker host and port fields.
  • Cached host/port strings.
  • An empty StorageInfo.
  • An empty RoaringBitmap.

Most of this metadata is identical across locations or unused for the majority of their lifetime. The eager allocations therefore contribute significant heap usage and GC pressure on the Master, Worker, and client.

The JOL benchmark includes both the live location graph and the weak interner overhead:

Scenario Before After, including interner Reduction
Allocator, 10,000 pairs 4,383,504 bytes 2,148,056 bytes 51.00%
Allocator, 1,000,000 pairs 424,143,504 bytes 184,308,056 bytes 56.55%
Packed decoded pair shape, 10,000 pairs 3,752,304 bytes 2,148,056 bytes 42.75%
Packed decoded pair shape, 1,000,000 pairs 375,200,304 bytes 184,308,056 bytes 50.88%

The additional serialization and range-read fixes are needed to ensure that lazy bitmap allocation does not introduce data loss, incorrect peer metadata, or compatibility regressions.

Does this PR resolve a correctness bug?

  • Yes

It fixes peer bitmap serialization and ensures that missing bitmap metadata does not cause an NPE or silently exclude a partition location during range reads.

Does this PR introduce any user-facing change?

  • Yes

There are no configuration, wire-protocol, or intended behavioral changes for users.

How was this patch tested?

The following verification was completed:

  • Ran the common, client, and master module regression suites:
    • 366 JUnit tests, with 1 ignored.
    • 325 ScalaTest tests.
    • 691 tests in total, with no failures or errors.
  • Ran the final PartitionLocationSuiteJ:
    • 20 tests passed.
    • Covers lazy initialization, endpoint interning, endpoint setters, ID parsing, peer references, compatibility getters, and same-version Java serialization.
  • Ran Worker partition and storage regression suites:
    • 27 Java tests passed.
    • 9 Scala tests passed.
  • Ran the Spark 3.5 LifecycleManagerReserveSlotsSuite integration test:
    • 1 test passed.
    • All 11 reactor modules completed successfully.
  • Compiled the affected modules through SBT:
    • celeborn-common/Test/compile
    • celeborn-service/Compile/compile
  • Ran spotless:check and git diff --check.
  • Manually ran PartitionLocationMemorySuiteJ with JOL for both the 10,000-pair and 1,000,000-pair scenarios.

Comment thread common/src/main/java/org/apache/celeborn/common/protocol/PartitionLocation.java Outdated
Comment thread common/pom.xml
@FMX

FMX commented Aug 25, 2026

Copy link
Copy Markdown
Contributor Author

Thanks for the review. I'll update this PR soon. @yew1eb @SteNicholas

@github-actions github-actions Bot added the correctness Correctness bugfix label Aug 25, 2026
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants