Updated client side metrics flow

Author: AmandeepSingh285

client/src/main/scala/org/apache/celeborn/client/ApplicationHeartbeater.scala

@@ -49,6 +49,8 @@ class ApplicationHeartbeater(
 
   private var stopped = false
   private val reviseLostShuffles = conf.reviseLostShufflesEnabled
+  private val appMetricLabels: util.Map[String, String] =
+    conf.clientMetricsAppLabels.asJava
 
   // Use independent app heartbeat threads to avoid being blocked by other operations.
   private val appHeartbeatIntervalMs = conf.appHeartbeatIntervalMs
@@ -89,7 +91,9 @@ class ApplicationHeartbeater(
                 workerStatusTracker.getNeedCheckedWorkers().toList.asJava,
                 ZERO_UUID,
                 true,
-                clientMetrics())
+                if (appMetricLabels.isEmpty) new util.HashMap[String, ClientMetric]()
+                else clientMetrics(),
+                appMetricLabels)
             val response = requestHeartbeat(appHeartbeat)
             if (response.statusCode == StatusCode.SUCCESS) {
               logDebug("Successfully send app heartbeat.")

client/src/main/scala/org/apache/celeborn/client/CelebornClientSource.scala

@@ -17,6 +17,8 @@
 
 package org.apache.celeborn.client
 
+import java.util.concurrent.ConcurrentHashMap
+
 import org.apache.celeborn.common.CelebornConf
 import org.apache.celeborn.common.metrics.{ClientMetric, MetricType}
 import org.apache.celeborn.common.metrics.source.{AbstractSource, Role}
@@ -29,6 +31,9 @@ class CelebornClientSource(conf: CelebornConf) extends AbstractSource(conf, Role
 
   import CelebornClientSource._
 
+  // Tracks previous counter values so we can send deltas to the master.
+  private val counterPrev = new ConcurrentHashMap[String, java.lang.Long]()
+
   addCounter(REGISTER_SHUFFLE_COUNT)
   addCounter(REGISTER_SHUFFLE_FAIL_COUNT)
   addCounter(UNREGISTER_SHUFFLE_COUNT)
@@ -39,34 +44,35 @@ class CelebornClientSource(conf: CelebornConf) extends AbstractSource(conf, Role
   addCounter(SHUFFLE_DATA_LOST_COUNT)
 
   def getMetricsSnapshot(): Map[String, ClientMetric] = {
-    val counterMetrics = counters().map(c =>
-      c.name -> ClientMetric(c.counter.getCount, MetricType.Counter))
+    // 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)
+      val delta = current - prev
+      if (delta > 0) Some(c.name -> ClientMetric(delta, MetricType.Counter))
+      else None
+    }
+    // Gauges: send the latest value as-is.
     val gaugeMetrics = gauges().map(g =>
       g.name -> ClientMetric(g.gauge.getValue.asInstanceOf[Number].longValue(), MetricType.Gauge))
     (counterMetrics ++ gaugeMetrics).toMap
   }
 
-  // start cleaner thread
-  startCleaner()
+  def start(): Unit = startCleaner()
+
+  def stop(): Unit = metricsCleaner.shutdown()
 }
 
 object CelebornClientSource {
-  // worker health
   val EXCLUDED_WORKER_COUNT = "ClientExcludedWorkerCount"
   val SHUTTING_WORKER_COUNT = "ClientShuttingWorkerCount"
-
-  // shuffle lifecycle
   val ACTIVE_SHUFFLE_COUNT = "ClientActiveShuffleCount"
   val REGISTER_SHUFFLE_COUNT = "ClientRegisterShuffleCount"
   val REGISTER_SHUFFLE_FAIL_COUNT = "ClientRegisterShuffleFailCount"
   val UNREGISTER_SHUFFLE_COUNT = "ClientUnregisterShuffleCount"
-
-  // write path
   val REVIVE_REQUEST_COUNT = "ClientReviveRequestCount"
   val REVIVE_FAIL_COUNT = "ClientReviveFailCount"
   val SLOT_RESERVATION_FAIL_COUNT = "ClientSlotReservationFailCount"
-
-  // data integrity
   val SHUFFLE_FETCH_FAILURE_COUNT = "ClientShuffleFetchFailureCount"
   val SHUFFLE_DATA_LOST_COUNT = "ClientShuffleDataLostCount"
 }

client/src/main/scala/org/apache/celeborn/client/ChangePartitionManager.scala

@@ -282,11 +282,9 @@ class ChangePartitionManager(
             None,
             lifecycleManager.workerStatusTracker.workerAvailableByLocation(req.oldPartition))))
       }
-      if (lifecycleManager.clientMetricsEnabled) {
-        lifecycleManager.clientSource.incCounter(
-          CelebornClientSource.REVIVE_FAIL_COUNT,
-          changePartitions.size)
-      }
+      lifecycleManager.incClientMetric(
+        CelebornClientSource.REVIVE_FAIL_COUNT,
+        changePartitions.size)
     }
 
     val candidates = new util.HashSet[WorkerInfo]()

client/src/main/scala/org/apache/celeborn/client/LifecycleManager.scala

@@ -223,20 +223,25 @@ class LifecycleManager(val appUniqueId: String, val conf: CelebornConf) extends
   }
 
   private val masterClient = new MasterClient(masterRpcEnvInUse, conf, false)
-  val clientSource = new CelebornClientSource(conf)
-  private[client] val clientMetricsEnabled = conf.metricsSystemEnable && conf.clientMetricsEnabled
+  private val clientMetricsEnabled = conf.metricsSystemEnable && conf.clientMetricsEnabled
+  val clientSource: Option[CelebornClientSource] =
+    if (clientMetricsEnabled) Some(new CelebornClientSource(conf)) else None
+
+  @inline private[client] def incClientMetric(name: String, delta: Long = 1L): Unit =
+    clientSource.foreach(_.incCounter(name, delta))
   val commitManager = new CommitManager(appUniqueId, conf, this)
   val workerStatusTracker = new WorkerStatusTracker(conf, this)
-  if (clientMetricsEnabled) {
-    clientSource.addGauge(CelebornClientSource.ACTIVE_SHUFFLE_COUNT) { () =>
+  clientSource.foreach { source =>
+    source.addGauge(CelebornClientSource.ACTIVE_SHUFFLE_COUNT) { () =>
       registeredShuffle.size
     }
-    clientSource.addGauge(CelebornClientSource.EXCLUDED_WORKER_COUNT) { () =>
+    source.addGauge(CelebornClientSource.EXCLUDED_WORKER_COUNT) { () =>
       workerStatusTracker.excludedWorkers.size
     }
-    clientSource.addGauge(CelebornClientSource.SHUTTING_WORKER_COUNT) { () =>
+    source.addGauge(CelebornClientSource.SHUTTING_WORKER_COUNT) { () =>
       workerStatusTracker.shuttingWorkers.size
     }
+    source.start()
   }
   private val heartbeater =
     new ApplicationHeartbeater(
@@ -252,11 +257,8 @@ class LifecycleManager(val appUniqueId: String, val conf: CelebornConf) extends
       registeredShuffle,
       reason => cancelAllActiveStages(reason),
       () =>
-        if (clientMetricsEnabled) {
-          clientSource.getMetricsSnapshot().asJava
-        } else {
-          new util.HashMap[String, ClientMetric]()
-        })
+        clientSource.map(_.getMetricsSnapshot().asJava)
+          .getOrElse(new util.HashMap[String, ClientMetric]()))
   private def resetFallbackCounts(counts: ConcurrentHashMap[String, java.lang.Long])
       : Map[String, java.lang.Long] = {
     val fallbackCounts = new util.HashMap[String, java.lang.Long]()
@@ -747,13 +749,9 @@ class LifecycleManager(val appUniqueId: String, val conf: CelebornConf) extends
 
     // Reply to all RegisterShuffle request for current shuffle id.
     def replyRegisterShuffle(response: RegisterShuffleResponse): Unit = {
-      if (clientMetricsEnabled) {
-        if (response.status == StatusCode.SUCCESS) {
-          clientSource.incCounter(CelebornClientSource.REGISTER_SHUFFLE_COUNT)
-        } else {
-          clientSource.incCounter(CelebornClientSource.REGISTER_SHUFFLE_FAIL_COUNT)
-        }
-      }
+      incClientMetric(
+        if (response.status == StatusCode.SUCCESS) CelebornClientSource.REGISTER_SHUFFLE_COUNT
+        else CelebornClientSource.REGISTER_SHUFFLE_FAIL_COUNT)
       registeringShuffleRequest.synchronized {
         val serializedMsg: Option[ByteBuffer] = partitionType match {
           case PartitionType.REDUCE =>
@@ -915,15 +913,11 @@ class LifecycleManager(val appUniqueId: String, val conf: CelebornConf) extends
       serdeVersion: SerdeVersion): Unit = {
     val contextWrapper =
       ChangeLocationsCallContext(context, partitionIds.size(), serdeVersion)
-    if (clientMetricsEnabled) {
-      clientSource.incCounter(CelebornClientSource.REVIVE_REQUEST_COUNT, partitionIds.size())
-    }
+    incClientMetric(CelebornClientSource.REVIVE_REQUEST_COUNT, partitionIds.size())
     // If shuffle not registered, reply ShuffleNotRegistered and return
     if (!registeredShuffle.contains(shuffleId)) {
       logError(s"[handleRevive] shuffle $shuffleId not registered!")
-      if (clientMetricsEnabled) {
-        clientSource.incCounter(CelebornClientSource.REVIVE_FAIL_COUNT)
-      }
+      incClientMetric(CelebornClientSource.REVIVE_FAIL_COUNT, partitionIds.size())
       contextWrapper.reply(
         -1,
         StatusCode.SHUFFLE_UNREGISTERED,
@@ -935,9 +929,7 @@ class LifecycleManager(val appUniqueId: String, val conf: CelebornConf) extends
       s"[handleRevive] shuffle $shuffleId, $mapIds, $partitionIds, $oldEpochs, $oldPartitions, $causes")
     if (commitManager.isStageEnd(shuffleId)) {
       logError(s"[handleRevive] shuffle $shuffleId stage ended!")
-      if (clientMetricsEnabled) {
-        clientSource.incCounter(CelebornClientSource.REVIVE_FAIL_COUNT)
-      }
+      incClientMetric(CelebornClientSource.REVIVE_FAIL_COUNT, partitionIds.size())
       contextWrapper.reply(
         -1,
         StatusCode.STAGE_ENDED,
@@ -1151,9 +1143,7 @@ class LifecycleManager(val appUniqueId: String, val conf: CelebornConf) extends
           if (invokeReportTaskShuffleFetchFailurePreCheck(taskId)) {
             logInfo(s"handle fetch failure for appShuffleId $appShuffleId shuffleId $shuffleId")
             ret = invokeAppShuffleTrackerCallback(appShuffleId)
-            if (ret && clientMetricsEnabled) {
-              clientSource.incCounter(CelebornClientSource.SHUFFLE_FETCH_FAILURE_COUNT)
-            }
+            if (ret) incClientMetric(CelebornClientSource.SHUFFLE_FETCH_FAILURE_COUNT)
             shuffleIds.put(appShuffleIdentifier, (shuffleId, false))
           } else {
             logInfo(
@@ -1285,9 +1275,7 @@ class LifecycleManager(val appUniqueId: String, val conf: CelebornConf) extends
           context.reply(MapperEndResponse(StatusCode.SUCCESS, serdeVersion))
         case false =>
           logError(s"Failed $message, reply ${StatusCode.SHUFFLE_DATA_LOST}.")
-          if (clientMetricsEnabled) {
-            clientSource.incCounter(CelebornClientSource.SHUFFLE_DATA_LOST_COUNT)
-          }
+          incClientMetric(CelebornClientSource.SHUFFLE_DATA_LOST_COUNT)
           context.reply(MapperEndResponse(StatusCode.SHUFFLE_DATA_LOST, serdeVersion))
       }
     }
@@ -1653,9 +1641,7 @@ class LifecycleManager(val appUniqueId: String, val conf: CelebornConf) extends
       // [[releasePartitionLocation]]. Now in the slots are all the successful partition
       // locations.
       logWarning(s"Reserve buffers for $shuffleId still fail after retrying, clear buffers.")
-      if (clientMetricsEnabled) {
-        clientSource.incCounter(CelebornClientSource.SLOT_RESERVATION_FAIL_COUNT)
-      }
+      incClientMetric(CelebornClientSource.SLOT_RESERVATION_FAIL_COUNT)
       destroySlotsWithRetry(shuffleId, slots)
     } else {
       logInfo(s"Reserve buffer success for shuffleId $shuffleId")
@@ -1875,9 +1861,7 @@ class LifecycleManager(val appUniqueId: String, val conf: CelebornConf) extends
           // if unregister shuffle not success, wait next turn
           if (StatusCode.SUCCESS == StatusCode.fromValue(unregisterShuffleResponse.getStatus)) {
             unregisterShuffleTime.remove(shuffleId)
-            if (clientMetricsEnabled) {
-              clientSource.incCounter(CelebornClientSource.UNREGISTER_SHUFFLE_COUNT)
-            }
+            incClientMetric(CelebornClientSource.UNREGISTER_SHUFFLE_COUNT)
           }
         }
       } else {
@@ -1889,9 +1873,7 @@ class LifecycleManager(val appUniqueId: String, val conf: CelebornConf) extends
         if (StatusCode.SUCCESS == StatusCode.fromValue(unregisterShuffleResponse.getStatus)) {
           shuffleIdsToRemove.foreach { shuffleId: Integer =>
             unregisterShuffleTime.remove(shuffleId)
-            if (clientMetricsEnabled) {
-              clientSource.incCounter(CelebornClientSource.UNREGISTER_SHUFFLE_COUNT)
-            }
+            incClientMetric(CelebornClientSource.UNREGISTER_SHUFFLE_COUNT)
           }
         }
       }
@@ -2101,6 +2083,7 @@ class LifecycleManager(val appUniqueId: String, val conf: CelebornConf) extends
    */
   override def stop(): Unit = {
     heartbeater.stop()
+    clientSource.foreach(_.stop())
     super.stop()
   }
 

client/src/main/scala/org/apache/celeborn/client/commit/ReducePartitionCommitHandler.scala

@@ -222,9 +222,7 @@ class ReducePartitionCommitHandler(
     } else {
       logError(s"Failed to handle stageEnd for $shuffleId, lost file!")
       dataLostShuffleSet.add(shuffleId)
-      if (conf.metricsSystemEnable && conf.clientMetricsEnabled) {
-        lifecycleManager.clientSource.incCounter(CelebornClientSource.SHUFFLE_DATA_LOST_COUNT)
-      }
+      lifecycleManager.incClientMetric(CelebornClientSource.SHUFFLE_DATA_LOST_COUNT)
       // record in stageEndShuffleSet
       setStageEnd(shuffleId)
     }

client/src/test/scala/org/apache/celeborn/client/WorkerStatusTrackerSuite.scala

@@ -167,7 +167,7 @@ class WorkerStatusTrackerSuite extends CelebornFunSuite {
     val lifecycleManager = new LifecycleManager("app-metrics-test", celebornConf)
     try {
       val statusTracker = lifecycleManager.workerStatusTracker
-      val source = lifecycleManager.clientSource
+      val source = lifecycleManager.clientSource.get
 
       val failed = new ShuffleFailedWorkers()
       val now = System.currentTimeMillis()

common/src/main/proto/TransportMessages.proto

@@ -502,6 +502,7 @@ message PbHeartbeatFromApplication {
   map<string, int64> applicationFallbackCounts = 10;
   PbHeartbeatInfo heartbeatInfo = 11;
   map<string, PbClientMetric> clientMetrics = 12;
+  map<string, string> metricLabels = 13;
 }
 
 enum PbMetricType {

common/src/main/scala/org/apache/celeborn/common/CelebornConf.scala

@@ -911,12 +911,17 @@ class CelebornConf(loadDefaults: Boolean) extends Cloneable with Logging with Se
   def metricsConf: Option[String] = get(METRICS_CONF)
   def metricsSystemEnable: Boolean = get(METRICS_ENABLED)
   def clientMetricsEnabled: Boolean = get(CLIENT_METRICS_ENABLED)
+  def masterClientMetricsEnabled: Boolean = get(MASTER_CLIENT_METRICS_ENABLED)
+  def masterClientMetricsRemovedAppRetentionMs: Long =
+    get(MASTER_CLIENT_METRICS_REMOVED_APP_RETENTION)
   def metricsSampleRate: Double = get(METRICS_SAMPLE_RATE)
   def metricsSlidingWindowSize: Int = get(METRICS_SLIDING_WINDOW_SIZE)
   def metricsCollectCriticalEnabled: Boolean = get(METRICS_COLLECT_CRITICAL_ENABLED)
   def metricsCapacity: Int = get(METRICS_CAPACITY)
   def metricsExtraLabels: Map[String, String] =
     get(METRICS_EXTRA_LABELS).map(Utils.parseKeyValuePair).toMap
+  def clientMetricsAppLabels: Map[String, String] =
+    get(CLIENT_METRICS_APP_LABELS).map(Utils.parseKeyValuePair).toMap
   def metricsWorkerAppTopResourceConsumptionCount: Int =
     get(METRICS_WORKER_APP_TOP_RESOURCE_CONSUMPTION_COUNT)
   def metricsWorkerAppTopResourceConsumptionBytesWrittenThreshold: Long =
@@ -5919,6 +5924,25 @@ object CelebornConf extends Logging {
       .booleanConf
       .createWithDefault(false)
 
+  val MASTER_CLIENT_METRICS_ENABLED: ConfigEntry[Boolean] =
+    buildConf("celeborn.metrics.master.clientMetrics.enabled")
+      .categories("metrics")
+      .doc("When true, the master exposes client-side metrics forwarded in application " +
+        "heartbeats on its Prometheus endpoint.")
+      .version("0.7.0")
+      .booleanConf
+      .createWithDefault(false)
+
+  val MASTER_CLIENT_METRICS_REMOVED_APP_RETENTION: ConfigEntry[Long] =
+    buildConf("celeborn.metrics.master.clientMetrics.removedApp.retentionMs")
+      .categories("metrics")
+      .doc("How long to retain removed application IDs in the client metrics source to " +
+        "reject late heartbeats after an application is lost. Entries older than this are " +
+        "periodically evicted.")
+      .version("0.7.0")
+      .timeConf(TimeUnit.MILLISECONDS)
+      .createWithDefaultString("5min")
+
   val METRICS_SAMPLE_RATE: ConfigEntry[Double] =
     buildConf("celeborn.metrics.sample.rate")
       .categories("metrics")
@@ -5965,6 +5989,20 @@ object CelebornConf extends Logging {
         "Allowed pattern is: `<label1_key>=<label1_value>[,<label2_key>=<label2_value>]*`")
       .createWithDefault(Seq.empty)
 
+  val CLIENT_METRICS_APP_LABELS: ConfigEntry[Seq[String]] =
+    buildConf("celeborn.client.metrics.appLabels")
+      .categories("client", "metrics")
+      .doc("Custom metric labels sent from the client in each application heartbeat and applied " +
+        "to client metrics exposed on the master's Prometheus endpoint. " +
+        "Labels' pattern is: `<label1_key>=<label1_value>[,<label2_key>=<label2_value>]*`; e.g. `env=prod,version=1`")
+      .version("0.7.0")
+      .stringConf
+      .toSequence
+      .checkValue(
+        labels => labels.map(_ => Try(Utils.parseKeyValuePair(_))).forall(_.isSuccess),
+        "Allowed pattern is: `<label1_key>=<label1_value>[,<label2_key>=<label2_value>]*`")
+      .createWithDefault(Seq.empty)
+
   val METRICS_WORKER_APP_TOP_RESOURCE_CONSUMPTION_COUNT: ConfigEntry[Int] =
     buildConf("celeborn.metrics.worker.app.topResourceConsumption.count")
       .categories("metrics")