Skip to content

[CELEBORN-194] Introduce client side metrics for celeborn#3740

Open
AmandeepSingh285 wants to merge 7 commits into
apache:mainfrom
AmandeepSingh285:adding-client-side-metrics
Open

[CELEBORN-194] Introduce client side metrics for celeborn#3740
AmandeepSingh285 wants to merge 7 commits into
apache:mainfrom
AmandeepSingh285:adding-client-side-metrics

Conversation

@AmandeepSingh285

@AmandeepSingh285 AmandeepSingh285 commented Jun 15, 2026

Copy link
Copy Markdown
Contributor

What changes were proposed in this pull request?

Adding client side metrics for Celeborn via heartbeat to master.

Why are the changes needed?

These changes help increase observability for Celeborn clients.

Does this PR resolve a correctness bug?

  • Yes

Does this PR introduce any user-facing change?

  • Yes

How was this patch tested?

@AmandeepSingh285 AmandeepSingh285 changed the title [CELEBORN-194] WIP Adding client side metrics for celeborn [CELEBORN-194] [WIP] Adding client side metrics for celeborn Jun 15, 2026
@AmandeepSingh285

Copy link
Copy Markdown
Contributor Author

Hi @SteNicholas , @RexXiong could you please help with a high level review on the implementation design for change adding client side metrics.
Thanks!

@codecov

codecov Bot commented Jun 17, 2026

Copy link
Copy Markdown

Codecov Report

❌ Patch coverage is 81.86813% with 33 lines in your changes missing coverage. Please review.
✅ Project coverage is 58.77%. Comparing base (17159eb) to head (fe1aa14).
⚠️ Report is 1 commits behind head on main.

Files with missing lines Patch % Lines
.../org/apache/celeborn/client/LifecycleManager.scala 41.38% 15 Missing and 2 partials ⚠️
...born/common/protocol/message/ControlMessages.scala 76.00% 2 Missing and 4 partials ⚠️
...pache/celeborn/client/ChangePartitionManager.scala 0.00% 3 Missing ⚠️
...pache/celeborn/client/ApplicationHeartbeater.scala 77.78% 1 Missing and 1 partial ⚠️
.../apache/celeborn/common/metrics/ClientMetric.scala 50.00% 1 Missing and 1 partial ⚠️
...eleborn/common/metrics/source/AbstractSource.scala 95.66% 1 Missing and 1 partial ⚠️
...n/client/commit/ReducePartitionCommitHandler.scala 0.00% 1 Missing ⚠️
Additional details and impacted files
@@             Coverage Diff              @@
##               main    #3740      +/-   ##
============================================
+ Coverage     57.73%   58.77%   +1.04%     
- Complexity      214      319     +105     
============================================
  Files           397      399       +2     
  Lines         27880    28056     +176     
  Branches       2714     2729      +15     
============================================
+ Hits          16095    16488     +393     
+ Misses        10635    10384     -251     
- Partials       1150     1184      +34     

☔ 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.

@AmandeepSingh285

Copy link
Copy Markdown
Contributor Author

Gentle ping @SteNicholas , @RexXiong could you please help with a high level review of the approach. Thanks!

@AmandeepSingh285 AmandeepSingh285 changed the title [CELEBORN-194] [WIP] Adding client side metrics for celeborn [CELEBORN-194] Adding client side metrics for celeborn Jun 22, 2026
@AmandeepSingh285

Copy link
Copy Markdown
Contributor Author

Gentle follow-up ping @SteNicholas @RexXiong . Would appreciate a high-level review of the proposed approach whenever you have some time. Thanks!

@AmandeepSingh285 AmandeepSingh285 changed the title [CELEBORN-194] Adding client side metrics for celeborn [CELEBORN-194][WIP] Adding client side metrics for celeborn Jun 22, 2026
@SteNicholas SteNicholas requested a review from Copilot June 22, 2026 10:13

Copilot AI left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Pull request overview

Note

Copilot couldn't run its full agentic review because no GitHub Actions runner was available. Make sure your repository has a runner available to run Copilot's review, or add a copilot-setup-steps.yml file specifying one with the runs-on attribute. See the docs for more details.

Adds client-side metrics collection in Celeborn clients and ships those metrics to the master via application heartbeats, where they are re-exposed on the master Prometheus endpoint labeled by applicationId.

Changes:

  • Extend HeartbeatFromApplication (and protobuf serde) to carry a clientMetrics map of {name -> (value, type)}.
  • Add client and master metric sources (CelebornClientSource, ApplicationMetricsSource) plus wiring in LifecycleManager/Master.
  • Introduce celeborn.client.metrics.enabled config and add/unit-test coverage for serde + source behavior.

Reviewed changes

Copilot reviewed 17 out of 17 changed files in this pull request and generated 6 comments.

Show a summary per file
File Description
master/src/test/scala/org/apache/celeborn/service/deploy/master/ApplicationMetricsSourceSuite.scala Adds unit tests for master-side application metrics source behavior.
master/src/main/scala/org/apache/celeborn/service/deploy/master/Master.scala Registers the new application metrics source and plumbs heartbeat clientMetrics through.
master/src/main/scala/org/apache/celeborn/service/deploy/master/ApplicationMetricsSource.scala Implements master-side cache + Prometheus re-export of client metrics by applicationId.
docs/configuration/metrics.md Documents new celeborn.client.metrics.enabled config.
common/src/test/scala/org/apache/celeborn/common/util/UtilsSuite.scala Adds serde round-trip test for clientMetrics in heartbeats.
common/src/main/scala/org/apache/celeborn/common/protocol/message/ControlMessages.scala Extends heartbeat message, protobuf encoding/decoding for client metrics.
common/src/main/scala/org/apache/celeborn/common/metrics/source/AbstractSource.scala Adds Role.CLIENT label behavior and counterExists helper.
common/src/main/scala/org/apache/celeborn/common/metrics/ClientMetric.scala Introduces ClientMetric + MetricType shared representation.
common/src/main/scala/org/apache/celeborn/common/CelebornConf.scala Adds celeborn.client.metrics.enabled config entry and accessor.
common/src/main/proto/TransportMessages.proto Adds clientMetrics field + metric type/message definitions to heartbeat protobuf.
client/src/test/scala/org/apache/celeborn/client/WorkerStatusTrackerSuite.scala Adds test coverage for excluded-worker metrics behavior.
client/src/test/scala/org/apache/celeborn/client/CelebornClientSourceSuite.scala Adds unit tests for client metric source counters/gauges + snapshot types.
client/src/main/scala/org/apache/celeborn/client/commit/ReducePartitionCommitHandler.scala Increments client “shuffle data lost” metric on lost-file conditions when enabled.
client/src/main/scala/org/apache/celeborn/client/LifecycleManager.scala Creates client metrics source, registers gauges, increments counters, and supplies snapshots to heartbeats.
client/src/main/scala/org/apache/celeborn/client/ChangePartitionManager.scala Increments revive-failure metrics when change partition assignment fails.
client/src/main/scala/org/apache/celeborn/client/CelebornClientSource.scala Implements client-side metrics source + snapshot export for heartbeat payload.
client/src/main/scala/org/apache/celeborn/client/ApplicationHeartbeater.scala Adds callback to attach client metrics to each HeartbeatFromApplication.

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

Comment thread client/src/main/scala/org/apache/celeborn/client/LifecycleManager.scala Outdated
Comment thread client/src/main/scala/org/apache/celeborn/client/CelebornClientSource.scala Outdated
Comment thread client/src/test/scala/org/apache/celeborn/client/WorkerStatusTrackerSuite.scala Outdated

@SteNicholas SteNicholas left a comment

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

@AmandeepSingh285, thanks for working on client-side metrics — the overall shape (client AbstractSource snapshot → heartbeat → master re-expose) is reasonable and the serde/config/docs are wired up. Since it's marked [WIP], I'm leaving review comments rather than approving; there are a few correctness/lifecycle issues worth resolving first (inline). Summary, most-impactful first:

  1. Client metrics source + cleaner thread leak (per Spark driver). clientSource is created unconditionally and never destroyed — see inline on LifecycleManager.scala:226.
  2. Master re-registers dead apps' metrics → permanent leak. Master is a plain RpcEndpoint (not ThreadSafeRpcEndpoint), so its Inbox runs with enableConcurrent = true and heartbeats are processed concurrently. A heartbeat racing/after handleAppLost resurrects the app's per-app gauges/counters, which are then never cleaned. See inline on ApplicationMetricsSource.updateApplicationMetrics.
  3. Non-atomic counter delta + fragile absolute→delta conversion. Concurrent heartbeats for one app double-count; app-restart-with-same-id or heartbeat reordering corrupt the delta. See inline on updateCounter.
  4. Unbounded master cardinality + silent truncation. Per-applicationId labeling has no top-N cap and no master-side enable flag; AbstractSource.getMetrics truncates at metricsCapacity (4096) and emits counters last, so app counters drop first. The codebase already solved this for worker per-app metrics via celeborn.metrics.worker.app.topResourceConsumption.count (default 0/off). See inline on Master.scala.
  5. Gauge flaps to 0 on removal (cache cleared before the gauge is unregistered) — inline on removeApplicationMetrics.
  6. Metric semantics: ClientReviveFailCount is incremented by changePartitions.size in one place but by +1 in handleRevive, and ClientShuffleDataLostCount is bumped in both handleMapPartitionEnd and ReducePartitionCommitHandler.stageEnd — mixed units / possible double-count. Inline on ChangePartitionManager.
  7. fromPb silently maps unknown PbMetricType to Gauge (forward-compat trap) — inline on ControlMessages.scala.

Minor / cleanup (no inline needed): the if (clientMetricsEnabled) clientSource.incCounter(...) guard is copy-pasted ~12×, and ReducePartitionCommitHandler recomputes the gate inline instead of reusing the cached clientMetricsEnabled field — a single incClientMetric(name, n) helper would centralize the gate and avoid the semantic drift in (6). ClientMetric/MetricType also duplicate proto PbClientMetric/PbMetricType (4 spots to keep in lockstep). Test gap: the counter-delta path in ApplicationMetricsSource is untested (ApplicationMetricsSourceSuite only sends MetricType.Gauge).

Comment thread client/src/main/scala/org/apache/celeborn/client/LifecycleManager.scala Outdated
Comment thread client/src/main/scala/org/apache/celeborn/client/ChangePartitionManager.scala Outdated
@AmandeepSingh285

Copy link
Copy Markdown
Contributor Author

Thanks @SteNicholas for the review. Still working on improving this PR. Will take into account all the updated you mentioned. Thanks!

@AmandeepSingh285 AmandeepSingh285 changed the title [CELEBORN-194][WIP] Adding client side metrics for celeborn [CELEBORN-194] Adding client side metrics for celeborn Jul 2, 2026
@SteNicholas SteNicholas changed the title [CELEBORN-194] Adding client side metrics for celeborn [CELEBORN-194] Introduce client side metrics for celeborn Jul 2, 2026
@SteNicholas SteNicholas requested a review from Copilot July 2, 2026 08:30
@SteNicholas

Copy link
Copy Markdown
Member

@AmandeepSingh285, please firstly resolve conflicts.

Copilot AI left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Pull request overview

Copilot reviewed 19 out of 19 changed files in this pull request and generated 3 comments.

Comment thread client/src/test/scala/org/apache/celeborn/client/WorkerStatusTrackerSuite.scala Outdated
@AmandeepSingh285 AmandeepSingh285 force-pushed the adding-client-side-metrics branch from 36a9603 to 438f541 Compare July 2, 2026 09:11

Copilot AI left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Pull request overview

Copilot reviewed 19 out of 19 changed files in this pull request and generated 6 comments.

Comment thread client/src/test/scala/org/apache/celeborn/client/WorkerStatusTrackerSuite.scala Outdated
Comment thread client/src/main/scala/org/apache/celeborn/client/ApplicationHeartbeater.scala Outdated
Comment thread client/src/main/scala/org/apache/celeborn/client/LifecycleManager.scala Outdated

Copilot AI left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Pull request overview

Copilot reviewed 19 out of 19 changed files in this pull request and generated 2 comments.

Comment on lines +6057 to +6059
.checkValue(
labels => labels.map(_ => Try(Utils.parseKeyValuePair(_))).forall(_.isSuccess),
"Allowed pattern is: `<label1_key>=<label1_value>[,<label2_key>=<label2_value>]*`")
Comment on lines +68 to +70
if (!masterClientMetricsEnabled || metricLabels.isEmpty || removedAppIds.containsKey(appId)) {
return
}

@SteNicholas SteNicholas left a comment

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

@AmandeepSingh285, thanks for updates. The proto evolution, Scala serde, pattern-match exhaustiveness, and master forwarding are all correct, and the test coverage is broad. I traced the whole path (client CelebornClientSource → heartbeat serde → master ApplicationMetricsSource / AbstractSource) and have a few items I'd like resolved before approving; leaving inline comments rather than approving.

The main theme: the master-side per-app aggregation in AbstractSource/ApplicationMetricsSource mutates several maps (removedAppIds, namedGaugesWithDetails/namedCountersWithDetails, the per-metric additionalDetails set, namedCounters, and the registry) without a unifying lock — and the Master endpoint is a plain RpcEndpoint, not ThreadSafeRpcEndpoint, so Inbox sets enableConcurrent = true and HeartbeatFromApplication runs concurrently with CheckForApplicationTimeOut/ApplicationLost. That makes the races below reachable, not theoretical. Secondary items: a silent-no-op enablement path, a client delta-loss on failed heartbeats, and a few test gaps. Details inline.

addCounter(name, labels)
AppMetricDetails(
name,
namedCounters.get(metricKey).counter,

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

NPE under the app-lost/heartbeat race. Inside computeIfAbsent, this does addCounter(name, labels) and then namedCounters.get(metricKey).counter, assuming the just-added counter is still present. But namedCounters is a different map, not covered by the namedCountersWithDetails bin lock. Because the Master endpoint dispatches concurrently (see the review summary), a concurrent removeAppFromTracked deregister for the same label key can run removeCounter(metricKey) (→ namedCounters.remove) in the window between addCounter's putIfAbsent (a no-op if the counter already existed) and this .get, making namedCounters.get(metricKey) return null → NPE thrown out of updateApplicationMetrics into handleHeartbeatFromApplication.

Consider having addCounter return the NamedCounter (or holding a single lock around the tracked-metric mutations) so this never observes a null.

while (iter.hasNext) {
val details = iter.next().getValue
details.additionalDetails.remove(appId)
if (details.additionalDetails.isEmpty) {

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Lost increment + transient deregistration under concurrency. For two apps sharing one label set, removeAppFromTracked (additionalDetails.remove(appId)isEmptyiter.remove()deregister) is not synchronized with addOrUpdateCounterForApp/addOrUpdateGaugeForApp (computeIfAbsentadditionalDetails.add(appId)handle.inc/set).

Interleaving: key K has additionalDetails = {app-1}. Thread A (heartbeat, app-2's first contribution to K) gets the existing details from computeIfAbsent. Thread B (app-1 lost) sees the set as {app-1} → empty → iter.remove() + removeCounter(K). Thread A resumes: additionalDetails.add("app-2") + handle.inc(delta) on the now-orphaned handle. Result: app-2's delta is permanently lost and the metric disappears from /metrics until app-2's next heartbeat recreates it from 0. Since Master is a concurrent endpoint and app-lost fires from timeoutDeadApplications, this window is real.

(Minor, same struct: AppMetricDetails.additionalDetails actually holds contributing appIds — a name like appIds/contributingAppIds would read far more clearly, since this set is the refcount that drives deregistration.)

appId: String,
metricLabels: Map[String, String],
metrics: JMap[String, ClientMetric]): Unit = {
if (!masterClientMetricsEnabled || metricLabels.isEmpty || removedAppIds.containsKey(appId)) {

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

TOCTOU between this gate and removeApplicationMetrics → resurrected, permanently-leaked metrics. This checks removedAppIds.containsKey(appId) and then (in addOrUpdate*ForApp) mutates the metric maps in separate, non-atomic steps relative to removeApplicationMetrics (which does removedAppIds.put + removeAppFromMetrics). On the concurrent Master endpoint:

  1. Thread A (heartbeat for app-1) reads containsKey("app-1") == false.
  2. Thread B (timeoutDeadApplicationshandleApplicationLost) runs removeApplicationMetrics("app-1") fully.
  3. Thread A resumes and re-registers app-1 into additionalDetails and the registry.

app-1 is now gated in removedAppIds (so removeApplicationMetrics won't run again) yet its metric series stays registered; after the retention cleaner evicts app-1 from removedAppIds, nothing ever deregisters that series → a dead app's series leaks indefinitely.

Relatedly, series are only reclaimed via handleAppLost; on a master failover removedAppIds (in-memory) is lost, and with high-cardinality operator appLabels there's no upper bound on distinct series. Worth a cap/warn and/or documenting that labels must be low-cardinality.

ZERO_UUID,
true)
true,
if (appMetricLabels.isEmpty) java.util.Collections.emptyMap[String, ClientMetric]()

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Silent no-op unless celeborn.client.metrics.appLabels is set. The client only sends metrics when appMetricLabels is non-empty, and the master's updateApplicationMetrics early-returns when metricLabels.isEmpty. So an operator who turns on the two obvious switches — celeborn.client.metrics.enabled=true and celeborn.metrics.master.clientMetrics.enabled=true — but leaves appLabels at its default (empty) gets zero metrics and zero log output. Neither enable-flag's doc mentions that appLabels is mandatory.

At minimum, log a one-time warn on the client when metrics are enabled but appLabels is empty. Better: don't gate emission on labels at all — emit with just role/instance so the default enablement path produces something.

// Counters: compute delta since last snapshot
val counterMetrics = counters().flatMap { c =>
val current = c.counter.getCount
val prev = Option(counterPrev.put(c.name, current)).map(_.longValue()).getOrElse(0L)

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Counter deltas are lost when a heartbeat fails to send. counterPrev.put(c.name, current) advances prev to current while the heartbeat message is being built; if requestHeartbeat(...) then fails (master unreachable/timeout) there's no retry and no rollback, so the emitted delta is discarded. The master's counter permanently under-counts by every failed heartbeat's delta.

(The symmetric hazard — advancing only on success — would double-count the timeout-but-actually-applied case, so a fully correct fix needs an ack/idempotency key. At least worth a comment acknowledging the at-most-once semantics, or carrying-forward the un-acked delta into the next snapshot.)

AppMetricDetails(name, holder, labels, ConcurrentHashMap.newKeySet[String]())
})
details.additionalDetails.add(appId)
details.handle.set(value)

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Gauge aggregation is last-writer-wins across apps sharing a label set (counters sum, gauges clobber). appId is intentionally kept only in additionalDetails, never in the metric key, so N apps that share a label set (e.g. coarse env=prod) collapse to one series and this handle.set(value) makes the exported gauge flap to whichever app heartbeated last — not a sum or max. A dashboard reading ClientActiveShuffleCount then sees an arbitrary single app's value. The PR's own test asserts this, so it looks intended, but the counter-vs-gauge asymmetry is surprising; please at least document it (and consider max/sum semantics, or requiring a per-app label for gauges).

Assert.assertEquals(2L, snapshot(CelebornClientSource.EXCLUDED_WORKER_COUNT).value)

// re-recording already-excluded workers does not change the gauge
statusTracker.recordWorkerFailure(failed)

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Dead comment / missing assertion, and this is an integration test in a unit suite. The // re-recording already-excluded workers does not change the gauge behavior is never checked — there's no assertEquals after this second recordWorkerFailure(failed), so a regression that double-counted (gauge → 4) would still pass. Add the assertion.

Separately, new LifecycleManager("app-metrics-test", celebornConf) (line 167) is heavy for a unit test: the constructor creates an RpcEnv (binds an ephemeral socket), runs initialize()registerApplicationInfo() against a non-existent master, and starts the heartbeater/commit/changePartition/release background threads that then spin retrying connections. A bare WorkerStatusTracker + CelebornClientSource (wired to the excluded-workers gauge) would cover this gauge deterministically without the port bind and connection noise.

assert(metrics.contains("""metrics_ClientRegisterShuffleCount_Count"""))
assert(metrics.contains("""role="Client""""))

val snapshot = source.getMetricsSnapshot()

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The delta semantics — the whole point of counterPrev — is never exercised. Every test calls getMetricsSnapshot() only once, so it only ever sees the first (full-value) delta. If the delta code regressed to emit cumulative counts (dropping counterPrev) or to keep emitting unchanged counters, the master would double-count every heartbeat, yet all three tests here would still pass. Please add a test that increments, snapshots, increments again, snapshots again, and asserts the second snapshot returns only the new delta and omits unchanged counters.

val heartbeatTrans = Utils.fromTransportMessage(Utils.toTransportMessage(heartbeat))
.asInstanceOf[HeartbeatFromApplication]

assert(heartbeatTrans.clientMetrics == clientMetrics)

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The roundtrip assertion is sound for clientMetrics (Java HashMap.equals + ClientMetric/MetricType case-class equality), but metricLabels (new proto field 13) is left at the default empty map and never round-tripped. A bug that dropped or mis-copied metricLabels on the wire would ship undetected. Construct the heartbeat with a non-empty metricLabels and assert it survives serde too. (Also worth noting: the master-side ApplicationMetricsSourceSuite creates ~13 enabled sources whose metricsCleaner daemon thread is never shut down — an afterEach/destroy() would avoid leaking scheduled threads across the suite.)

.createWithDefault(false)

val MASTER_CLIENT_METRICS_REMOVED_APP_RETENTION: ConfigEntry[Long] =
buildConf("celeborn.metrics.master.clientMetrics.removedApp.retentionMs")

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Minor config-naming nit: celeborn.metrics.master.clientMetrics.removedApp.retentionMs embeds the unit (Ms) in the key of a timeConf, which reads self-contradictorily against its own default (retentionMs = 5min) and invites passing a raw millisecond number. No other timeConf in the repo suffixes its unit (siblings use .timeout, .interval, .expireTimeout, .retention). Suggest ...removedApp.retention (or .expireTimeout). While here: CLIENT_METRICS_ENABLED uses .categories("metrics") but its sibling CLIENT_METRICS_APP_LABELS uses .categories("client", "metrics"), so the enable flag is missing from docs/configuration/client.md — consider .categories("client", "metrics") for consistency.

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