SteNicholas commented on PR #3672:
URL: https://github.com/apache/celeborn/pull/3672#issuecomment-4490826404
@xumingming, please take a look at the comments from claude code:
```
Code Review — [CELEBORN-2315] Add iterator fully-consumed validation after
shuffle write
(apache/celeborn#3672)
Overview
Adds a post-write safety check to HashBasedShuffleWriter and
SortBasedShuffleWriter (both spark-2
and spark-3): doWrite() now returns whether the input iterator still has
records, and
close(iteratorHasNext) calls SparkUtils.assertIteratorFullyConsumed(...)
which kills the task with
TaskKilledException if records remain.
TaskInterruptedHelper.throwTaskKillException is extended
to carry an optional message merged with Spark's kill reason. +171/-46
across 9 files.
Correctness — solid, with points to verify
- Check placement is well-reasoned and correct. In every writer the
assertion is placed after the
pusher is drained / buffers returned and before
shuffleClient.mapperEnd(...). This is the key
invariant — a partial map output is never committed to the shuffle service
— and it's right.
- mapSideCombine handling is a genuine improvement. Capturing
combinedIterator and checking
combinedIterator.hasNext() (rather than the raw records) is the correct
level: write0 consumes the
combined iterator, and its exhaustion implies the input was drained by
the aggregator.
- Null/reason merging in throwTaskKillException(String) covers all four
combinations cleanly.
- ⚠️ finally { if (needCleanupPusher) cleanupPusher(); } still runs when
the assertion throws.
TaskKilledException isn't an InterruptedException, so it bypasses the
catch; needCleanupPusher is
still true, so cleanupPusher() executes after the pusher already
terminated. Two things to
confirm: (1) cleanupPusher() is safe/idempotent post-termination, and (2)
if it throws
IOException, that exception will supersede the TaskKilledException (Java
finally semantics),
masking the real kill reason. Worth a comment or guarding.
- ⚠️ Generics: scala.collection.Iterator<?> combinedIterator =
dep.aggregator().get().combineValuesByKey(...) then
write0(combinedIterator). The original passed
the expression inline (type inferred). Confirm write0 accepts the wildcard
without an
unchecked/raw warning regression.
- Minor: records.hasNext() is assumed side-effect-free/idempotent. True
for standard Spark
iterators; fine to rely on but worth being aware of for exotic
interruptible/completion iterators.
Behavior change / risk
- This converts a silent condition into a hard task kill. For a
ShuffleMapTask the writer must
consume the full iterator, so a non-empty iterator genuinely indicates
data loss — semantically
correct. But any false positive (a custom iterator whose hasNext()
blocks/recomputes, or an
unforeseen early-termination path) would now fail previously-"successful"
tasks.
- Suggestion: gate the assertion behind a CelebornConf flag (default on)
so it can be disabled
quickly if false positives surface in production. Low cost, much safer
rollout for a correctness
guard touching every shuffle write.
Test coverage — too shallow for the risk
- The only new tests call
SparkUtils.assertIteratorFullyConsumed(false/true) directly. They do not
cover the actual integration: that a real writer with a
partially-consumed iterator triggers the
kill, that the normal full-consumption path does not falsely trip, that
the kill happens before
mapperEnd, or the new throwTaskKillException(String) branches (spark
reason + message).
- Recommend at least: (1) an assertion in the existing check(...) writer
suites that a normal full
write does not throw (guards against false positives across fast-write /
combine / plain paths);
(2) a test driving a writer/doWrite with an iterator left non-empty and
asserting
TaskKilledException is raised before mapperEnd.
Consistency / style
- @VisibleForTesting and the explanatory return-contract comment are
present only on spark-3
SortBasedShuffleWriter.doWrite; the new doWrite in spark-2 Sort and both
Hash writers have
neither. Make the four doWrite methods consistent (annotation + one-line
"returns true if not
fully consumed" doc).
- The "why this placement" comment exists on the Hash writers but not on
either
SortBasedShuffleWriter before assertIteratorFullyConsumed, even though
placement-before-mapperEnd
matters equally there. Add the same rationale so a future refactor doesn't
move it past mapperEnd.
- spark-3 test adds dependency.mapSideCombine()/aggregator() mocks but
spark-2 test does not —
confirm spark-2 suite still exercises the new doWrite branch (or is
relying on a real dependency).
Verdict
Sound, well-placed change with correct core logic and good handling of the
combine path. No
blocking bug found. Before merge I'd want: a config gate for safe rollout,
deeper
(integration-level) tests of the guard and the new message branches, the
finally/cleanupPusher
interaction confirmed, and the cross-writer @VisibleForTesting/comment
consistency tidied.
```
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
