From 27614330b1ecb9adeba13d44aa238b92827df033 Mon Sep 17 00:00:00 2001 From: "nilesh.rathi" Date: Tue, 4 Aug 2026 19:54:11 +0530 Subject: [PATCH] LIVY-866: Optimize Yarn GetApplications Query to prevent additional load on Yarn and Livy --- .../org/apache/livy/utils/SparkYarnApp.scala | 38 +++++++++++++------ .../apache/livy/utils/SparkYarnAppSpec.scala | 36 +++++++++++++++--- 2 files changed, 57 insertions(+), 17 deletions(-) diff --git a/server/src/main/scala/org/apache/livy/utils/SparkYarnApp.scala b/server/src/main/scala/org/apache/livy/utils/SparkYarnApp.scala index bd7d29fa5..358950842 100644 --- a/server/src/main/scala/org/apache/livy/utils/SparkYarnApp.scala +++ b/server/src/main/scala/org/apache/livy/utils/SparkYarnApp.scala @@ -16,6 +16,8 @@ */ package org.apache.livy.utils +import java.util + import scala.annotation.tailrec import scala.collection.JavaConverters._ import scala.collection.mutable.ArrayBuffer @@ -25,6 +27,7 @@ import scala.language.postfixOps import scala.util.Try import scala.util.control.NonFatal +import org.apache.hadoop.yarn.api.protocolrecords.GetApplicationsRequest import org.apache.hadoop.yarn.api.records.{ApplicationId, ApplicationReport, FinalApplicationStatus, YarnApplicationState} import org.apache.hadoop.yarn.client.api.YarnClient import org.apache.hadoop.yarn.conf.YarnConfiguration @@ -68,13 +71,23 @@ object SparkYarnApp extends Logging { private var sessionLeakageCheckInterval: Long = _ + /** + * Build a GetApplicationsRequest filtered by Spark application type and tags. + * Tags are normalized to lowercase to match YARN's tag storage behavior. + */ + private def createGetApplicationsRequest(appTags: util.Set[String]): GetApplicationsRequest = { + val normalizedTags = new util.HashSet[String]() + appTags.asScala.foreach(tag => normalizedTags.add(tag.toLowerCase)) + val request = GetApplicationsRequest.newInstance(appType) + request.setApplicationTags(normalizedTags) + request + } + private val leakedAppsGCThread = new Thread() { override def run(): Unit = { - val client = { - mockYarnClient match { - case Some(client) => client - case None => yarnClient - } + val client = mockYarnClient match { + case Some(client) => client + case None => yarnClient } while (true) { @@ -82,13 +95,16 @@ object SparkYarnApp extends Logging { // kill the app if found it and remove it if exceeding a threshold val iter = leakedAppTags.entrySet().iterator() val now = System.currentTimeMillis() - val apps = client.getApplications(appType).asScala + val tagSet = new util.HashSet[String](leakedAppTags.keySet()) + val request = createGetApplicationsRequest(tagSet) + val apps = client.getApplications(request).asScala while(iter.hasNext) { var isRemoved = false val entry = iter.next() + val tagLowerCase = entry.getKey.toLowerCase() - apps.find(_.getApplicationTags.contains(entry.getKey)) + apps.find(_.getApplicationTags.contains(tagLowerCase)) .foreach({ e => info(s"Kill leaked app ${e.getApplicationId}") client.killApplication(e.getApplicationId) @@ -196,10 +212,10 @@ class SparkYarnApp private[utils] ( } val appTagLowerCase = appTag.toLowerCase() - - // FIXME Should not loop thru all YARN applications but YarnClient doesn't offer an API. - // Consider calling rmClient in YarnClient directly. - yarnClient.getApplications(appType).asScala.find(_.getApplicationTags.contains(appTagLowerCase)) + val appTags: util.Set[String] = util.Collections.singleton(appTagLowerCase) + val request = createGetApplicationsRequest(appTags) + val applicationReports = yarnClient.getApplications(request) + applicationReports.asScala.find(_.getApplicationTags.contains(appTagLowerCase)) match { case Some(app) => app.getApplicationId case None => diff --git a/server/src/test/scala/org/apache/livy/utils/SparkYarnAppSpec.scala b/server/src/test/scala/org/apache/livy/utils/SparkYarnAppSpec.scala index 509b84609..1cc6b4b60 100644 --- a/server/src/test/scala/org/apache/livy/utils/SparkYarnAppSpec.scala +++ b/server/src/test/scala/org/apache/livy/utils/SparkYarnAppSpec.scala @@ -16,7 +16,6 @@ */ package org.apache.livy.utils -import java.util.ArrayList import java.util.concurrent.{CountDownLatch, TimeUnit} import java.util.concurrent.atomic.{AtomicBoolean, AtomicInteger} @@ -24,12 +23,15 @@ import scala.collection.JavaConverters._ import scala.concurrent.duration._ import scala.language.postfixOps +import org.apache.hadoop.yarn.api.protocolrecords.GetApplicationsRequest import org.apache.hadoop.yarn.api.records._ import org.apache.hadoop.yarn.api.records.FinalApplicationStatus.UNDEFINED import org.apache.hadoop.yarn.api.records.YarnApplicationState._ import org.apache.hadoop.yarn.client.api.YarnClient import org.apache.hadoop.yarn.exceptions.ApplicationAttemptNotFoundException import org.apache.hadoop.yarn.util.ConverterUtils +import org.mockito.ArgumentCaptor +import org.mockito.Matchers.any import org.mockito.Mockito._ import org.mockito.invocation.InvocationOnMock import org.mockito.stubbing.Answer @@ -49,6 +51,27 @@ class SparkYarnAppSpec extends FunSpec with LivyBaseUnitTestSuite { Thread.`yield`() } + private def mockGetApplicationsByTags( + client: YarnClient, + reports: List[ApplicationReport]): Unit = { + when(client.getApplications(any(classOf[GetApplicationsRequest]))) + .thenReturn(reports.asJava) + } + + private def verifyFilteredGetApplicationsRequest( + client: YarnClient, + expectedTags: Set[String]): Unit = { + val requestCaptor = ArgumentCaptor.forClass(classOf[GetApplicationsRequest]) + verify(client, atLeastOnce()).getApplications(requestCaptor.capture()) + val capturedRequest = requestCaptor.getValue + expectedTags.foreach { tag => + assert(capturedRequest.getApplicationTags.contains(tag.toLowerCase), + s"Request must contain the lowercase tag '$tag'") + } + assert(capturedRequest.getApplicationTypes.contains("SPARK"), + "Request must filter by application type 'SPARK'") + } + describe("SparkYarnApp") { val TEST_TIMEOUT = 30 seconds val appId = ConverterUtils.toApplicationId("application_1467912463905_0021") @@ -389,8 +412,7 @@ class SparkYarnAppSpec extends FunSpec with LivyBaseUnitTestSuite { when(mockAppReport.getFinalApplicationStatus).thenReturn(FinalApplicationStatus.SUCCEEDED) when(mockAppReport.getYarnApplicationState).thenReturn(YarnApplicationState.FINISHED) when(mockYarnClient.getApplicationReport(appId)).thenReturn(mockAppReport) - when(mockYarnClient.getApplications(Set("SPARK").asJava)) - .thenReturn(List(mockAppReport).asJava) + mockGetApplicationsByTags(mockYarnClient, List(mockAppReport)) val mockListener = mock[SparkAppListener] val mockSparkSubmit = mock[LineBufferedProcess] @@ -404,6 +426,7 @@ class SparkYarnAppSpec extends FunSpec with LivyBaseUnitTestSuite { verify(mockYarnClient, atLeast(1)).getApplicationReport(appId) verify(mockListener).appIdKnown(appId.toString) + verifyFilteredGetApplicationsRequest(mockYarnClient, Set(appTag)) } } } @@ -673,17 +696,18 @@ class SparkYarnAppSpec extends FunSpec with LivyBaseUnitTestSuite { livyConf.set(LivyConf.YARN_APP_LEAKAGE_CHECK_TIMEOUT, "1000ms") val client = mock[YarnClient] - when(client.getApplications(SparkYarnApp.appType)). - thenReturn(new ArrayList[ApplicationReport]()) + mockGetApplicationsByTags(client, List.empty) SparkYarnApp.init(livyConf, Some(client)) SparkYarnApp.leakedAppTags.clear() - SparkYarnApp.leakedAppTags.put("leakApp", System.currentTimeMillis()) + val leakAppTag = "leakApp" + SparkYarnApp.leakedAppTags.put(leakAppTag, System.currentTimeMillis()) Eventually.eventually(Eventually.timeout(TEST_TIMEOUT), Eventually.interval(100 millis)) { assert(SparkYarnApp.leakedAppTags.size() == 0) } + verifyFilteredGetApplicationsRequest(client, Set(leakAppTag)) } }