Skip to content

Native Celeborn shuffle: reflectively replacing Celeborn's final fields can release push admission while payloads are in flight #5528

Description

@andygrove

Following up on PR #5513 (part of #5352). CelebornTransportCallbackTracker installs its push-ownership instrumentation by reflectively assigning four fields on live Celeborn objects, and all four are declared private final in the Celeborn releases the PR targets.

The four are TransportClientFactory.clientBootstraps, TransportClient.channel, TransportResponseHandler.outstandingPushes, and ShuffleClientImpl.pushDataRetryPool. I checked the declarations in both v0.6.0 and v0.7.0 and they match in each:

TransportClientFactory.java:79   private final List<TransportClientBootstrap> clientBootstraps;
TransportClient.java:81          private final Channel channel;
TransportResponseHandler.java:60 private final ConcurrentHashMap<Long, PushRequestInfo> outstandingPushes;
ShuffleClientImpl.java:131       private final ExecutorService pushDataRetryPool;

setAccessible(true) makes the write legal for a non-static final field, so this is not an access problem. The issue is visibility. JLS 17.5.3 gives no guarantee that a final field modified after construction becomes visible to another thread unless that thread synchronizes with the writer. Comet writes under the transport factory monitor, the response handler monitor, or the shuffle client monitor, and Celeborn's own reads of these fields on the Netty event loops and on other tasks' pushing threads take none of those. A thread that has no happens-before edge, or a JIT that has already trusted and folded the final field load, keeps using the original channel, the original outstandingPushes map, and the original retry pool.

What concerns me is the direction the failure takes. If a push's callback is never wrapped because some thread still sees the original request map, then no Lease is retained for it, Push.owners reaches zero at close(), isComplete() returns true, and the pusher releases its admission bytes while the payload is still in flight. That is precisely the over-release the tracker was added to prevent, and it happens silently in the direction of memory overcommit rather than of a stall.

I do not think this is reachable by any test currently in the tree. The 46 tests in CelebornShufflePartitionPusherSuite all drive hand-written single-threaded Scala stand-ins, so there is never a concurrently running Celeborn reader racing the reflective write.

Is there a way to get the same push-completion signal without mutating another project's final fields? Two ideas that come to mind are wrapping the client object that Comet hands to the native writer so Comet controls the push entry point, or asking Celeborn upstream for a supported push-completion listener hook so this does not depend on private internals at all. If neither is workable soon, would it be worth documenting the reliance explicitly and pinning the supported Celeborn versions?

Activity

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Metadata

Metadata

Assignees

Labels

area:shuffleShuffle (JVM and native)bugSomething isn't workingcorrectnesspriority:mediumFunctional bugs, performance regressions, broken features

Type

No type

Projects

No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions