[CELEBORN-194] Introduce client side metrics for celeborn#3740
[CELEBORN-194] Introduce client side metrics for celeborn#3740AmandeepSingh285 wants to merge 7 commits into
Conversation
|
Hi @SteNicholas , @RexXiong could you please help with a high level review on the implementation design for change adding client side metrics. |
Codecov Report❌ Patch coverage is 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. 🚀 New features to boost your workflow:
|
|
Gentle ping @SteNicholas , @RexXiong could you please help with a high level review of the approach. Thanks! |
|
Gentle follow-up ping @SteNicholas @RexXiong . Would appreciate a high-level review of the proposed approach whenever you have some time. Thanks! |
There was a problem hiding this comment.
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 aclientMetricsmap of{name -> (value, type)}. - Add client and master metric sources (
CelebornClientSource,ApplicationMetricsSource) plus wiring inLifecycleManager/Master. - Introduce
celeborn.client.metrics.enabledconfig 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.
There was a problem hiding this comment.
@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:
- Client metrics source + cleaner thread leak (per Spark driver).
clientSourceis created unconditionally and never destroyed — see inline onLifecycleManager.scala:226. - Master re-registers dead apps' metrics → permanent leak.
Masteris a plainRpcEndpoint(notThreadSafeRpcEndpoint), so its Inbox runs withenableConcurrent = trueand heartbeats are processed concurrently. A heartbeat racing/afterhandleAppLostresurrects the app's per-app gauges/counters, which are then never cleaned. See inline onApplicationMetricsSource.updateApplicationMetrics. - 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. - Unbounded master cardinality + silent truncation. Per-
applicationIdlabeling has no top-N cap and no master-side enable flag;AbstractSource.getMetricstruncates atmetricsCapacity(4096) and emits counters last, so app counters drop first. The codebase already solved this for worker per-app metrics viaceleborn.metrics.worker.app.topResourceConsumption.count(default 0/off). See inline onMaster.scala. - Gauge flaps to 0 on removal (cache cleared before the gauge is unregistered) — inline on
removeApplicationMetrics. - Metric semantics:
ClientReviveFailCountis incremented bychangePartitions.sizein one place but by +1 inhandleRevive, andClientShuffleDataLostCountis bumped in bothhandleMapPartitionEndandReducePartitionCommitHandler.stageEnd— mixed units / possible double-count. Inline onChangePartitionManager. fromPbsilently maps unknownPbMetricTypeto Gauge (forward-compat trap) — inline onControlMessages.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).
|
Thanks @SteNicholas for the review. Still working on improving this PR. Will take into account all the updated you mentioned. Thanks! |
|
@AmandeepSingh285, please firstly resolve conflicts. |
36a9603 to
438f541
Compare
| .checkValue( | ||
| labels => labels.map(_ => Try(Utils.parseKeyValuePair(_))).forall(_.isSuccess), | ||
| "Allowed pattern is: `<label1_key>=<label1_value>[,<label2_key>=<label2_value>]*`") |
| if (!masterClientMetricsEnabled || metricLabels.isEmpty || removedAppIds.containsKey(appId)) { | ||
| return | ||
| } |
There was a problem hiding this comment.
@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, |
There was a problem hiding this comment.
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) { |
There was a problem hiding this comment.
Lost increment + transient deregistration under concurrency. For two apps sharing one label set, removeAppFromTracked (additionalDetails.remove(appId) → isEmpty → iter.remove() → deregister) is not synchronized with addOrUpdateCounterForApp/addOrUpdateGaugeForApp (computeIfAbsent → additionalDetails.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)) { |
There was a problem hiding this comment.
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:
- Thread A (heartbeat for
app-1) readscontainsKey("app-1") == false. - Thread B (
timeoutDeadApplications→handleApplicationLost) runsremoveApplicationMetrics("app-1")fully. - Thread A resumes and re-registers
app-1intoadditionalDetailsand 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]() |
There was a problem hiding this comment.
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) |
There was a problem hiding this comment.
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) |
There was a problem hiding this comment.
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) |
There was a problem hiding this comment.
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() |
There was a problem hiding this comment.
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) |
There was a problem hiding this comment.
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") |
There was a problem hiding this comment.
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.
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?
Does this PR introduce any user-facing change?
How was this patch tested?