From 990a6efbae2201916db1fa087a73e2a6600fd272 Mon Sep 17 00:00:00 2001 From: "nilesh.rathi" Date: Mon, 6 Jul 2026 11:22:06 +0530 Subject: [PATCH] [LIVY-1059] Add livy.spark.yarn.queue configuration to set default YARN queue --- conf/livy.conf.template | 4 ++ .../main/scala/org/apache/livy/LivyConf.scala | 4 ++ .../livy/server/batch/BatchSession.scala | 2 +- .../interactive/InteractiveSession.scala | 4 +- .../livy/server/batch/BatchSessionSpec.scala | 39 +++++++++++- .../interactive/InteractiveSessionSpec.scala | 62 +++++++++++++++++++ 6 files changed, 111 insertions(+), 4 deletions(-) diff --git a/conf/livy.conf.template b/conf/livy.conf.template index 0299dca7f..30a3a1f40 100644 --- a/conf/livy.conf.template +++ b/conf/livy.conf.template @@ -47,6 +47,10 @@ # What spark deploy mode Livy sessions should use. # livy.spark.deploy-mode = +# What default YARN queue Livy sessions should use if not specified in the client request. +# By default, it is null and falls back to the global Hadoop cluster default queue. +# livy.spark.yarn.queue = default + # Configure Livy server http request and response header size. # livy.server.request-header.size = 131072 # livy.server.response-header.size = 131072 diff --git a/server/src/main/scala/org/apache/livy/LivyConf.scala b/server/src/main/scala/org/apache/livy/LivyConf.scala index 006ab5dff..46c851f35 100644 --- a/server/src/main/scala/org/apache/livy/LivyConf.scala +++ b/server/src/main/scala/org/apache/livy/LivyConf.scala @@ -44,6 +44,7 @@ object LivyConf { val SPARK_HOME = Entry("livy.server.spark-home", null) val LIVY_SPARK_MASTER = Entry("livy.spark.master", "local") val LIVY_SPARK_DEPLOY_MODE = Entry("livy.spark.deploy-mode", null) + val SPARK_YARN_QUEUE = Entry("livy.spark.yarn.queue", null) // Two configurations to specify Spark and related Scala version. These are internal // configurations will be set by LivyServer and used in session creation. It is not required to @@ -485,6 +486,9 @@ class LivyConf(loadDefaults: Boolean) extends ClientConf[LivyConf](null) { /** Return the spark master Livy sessions should use. */ def sparkMaster(): String = get(LIVY_SPARK_MASTER) + /** Return the default YARN queue Livy sessions should use. */ + def sparkYarnQueue(): Option[String] = Option(get(SPARK_YARN_QUEUE)).filterNot(_.isEmpty) + /** Return the path to the spark-submit executable. */ def sparkSubmit(): String = { sparkHome().map { _ + File.separator + "bin" + File.separator + "spark-submit" }.get diff --git a/server/src/main/scala/org/apache/livy/server/batch/BatchSession.scala b/server/src/main/scala/org/apache/livy/server/batch/BatchSession.scala index 8b64a0398..ec0f961f4 100644 --- a/server/src/main/scala/org/apache/livy/server/batch/BatchSession.scala +++ b/server/src/main/scala/org/apache/livy/server/batch/BatchSession.scala @@ -83,7 +83,7 @@ object BatchSession extends Logging { request.executorMemory.foreach(builder.executorMemory) request.executorCores.foreach(builder.executorCores) request.numExecutors.foreach(builder.numExecutors) - request.queue.foreach(builder.queue) + request.queue.orElse(livyConf.sparkYarnQueue()).foreach(builder.queue) request.name.foreach(builder.name) sessionStore.save(BatchSession.RECOVERY_SESSION_TYPE, s.recoveryMetadata) diff --git a/server/src/main/scala/org/apache/livy/server/interactive/InteractiveSession.scala b/server/src/main/scala/org/apache/livy/server/interactive/InteractiveSession.scala index 0667b718c..0c7cb0bde 100644 --- a/server/src/main/scala/org/apache/livy/server/interactive/InteractiveSession.scala +++ b/server/src/main/scala/org/apache/livy/server/interactive/InteractiveSession.scala @@ -107,7 +107,7 @@ object InteractiveSession extends Logging { SparkLauncher.EXECUTOR_MEMORY -> request.executorMemory.map(_.toString), "spark.executor.instances" -> request.numExecutors.map(_.toString), "spark.app.name" -> request.name.map(_.toString), - "spark.yarn.queue" -> request.queue + "spark.yarn.queue" -> request.queue.orElse(livyConf.sparkYarnQueue()) ) userOpts.foreach { case (key, opt) => @@ -152,7 +152,7 @@ object InteractiveSession extends Logging { request.jars, request.numExecutors, request.pyFiles, - request.queue, + request.queue.orElse(livyConf.sparkYarnQueue()), mockApp) } diff --git a/server/src/test/scala/org/apache/livy/server/batch/BatchSessionSpec.scala b/server/src/test/scala/org/apache/livy/server/batch/BatchSessionSpec.scala index 401a8beb1..807b20f34 100644 --- a/server/src/test/scala/org/apache/livy/server/batch/BatchSessionSpec.scala +++ b/server/src/test/scala/org/apache/livy/server/batch/BatchSessionSpec.scala @@ -33,7 +33,7 @@ import org.apache.livy.{LivyBaseUnitTestSuite, LivyConf, Utils} import org.apache.livy.server.AccessManager import org.apache.livy.server.recovery.SessionStore import org.apache.livy.sessions.SessionState -import org.apache.livy.utils.{AppInfo, Clock, SparkApp} +import org.apache.livy.utils.{AppInfo, Clock, SparkApp, SparkProcessBuilder} class BatchSessionSpec extends FunSpec @@ -138,6 +138,43 @@ class BatchSessionSpec }) should be (true) } + it("should pass default YARN queue to spark-submit when request queue is empty") { + val req = new CreateBatchRequest() + req.file = script.toString + req.queue = None + req.conf = Map("spark.driver.extraClassPath" -> sys.props("java.class.path")) + + val conf = new LivyConf() + .set(LivyConf.LOCAL_FS_WHITELIST, sys.props("java.io.tmpdir")) + .set(LivyConf.SPARK_YARN_QUEUE, "livy-default-batch-queue") + + val accessManager = new AccessManager(conf) + val batch = BatchSession.create(10, None, req, conf, accessManager, null, None, sessionStore) + batch.start() + + Utils.waitUntil({ () => !batch.state.isActive }, Duration(10, TimeUnit.SECONDS)) + batch.logLines().mkString should include("livy-default-batch-queue") + } + + it("should prioritize user-specified request queue over LivyConf configuration") { + val req = new CreateBatchRequest() + req.file = script.toString + req.queue = Some("user-custom-batch-queue") + req.conf = Map("spark.driver.extraClassPath" -> sys.props("java.class.path")) + + val conf = new LivyConf() + .set(LivyConf.LOCAL_FS_WHITELIST, sys.props("java.io.tmpdir")) + .set(LivyConf.SPARK_YARN_QUEUE, "livy-default-batch-queue") + + val accessManager = new AccessManager(conf) + val batch = BatchSession.create(20, None, req, conf, accessManager, null, None, sessionStore) + batch.start() + + Utils.waitUntil({ () => !batch.state.isActive }, Duration(10, TimeUnit.SECONDS)) + batch.logLines().mkString should include("user-custom-batch-queue") + batch.logLines().mkString should not include "livy-default-batch-queue" + } + def testRecoverSession(name: Option[String]): Unit = { val conf = new LivyConf() val req = new CreateBatchRequest() diff --git a/server/src/test/scala/org/apache/livy/server/interactive/InteractiveSessionSpec.scala b/server/src/test/scala/org/apache/livy/server/interactive/InteractiveSessionSpec.scala index e7d651f89..31daefa1f 100644 --- a/server/src/test/scala/org/apache/livy/server/interactive/InteractiveSessionSpec.scala +++ b/server/src/test/scala/org/apache/livy/server/interactive/InteractiveSessionSpec.scala @@ -322,4 +322,66 @@ class InteractiveSessionSpec extends FunSpec s.logLines().mkString should include("RSCDriver URI is unknown") } } + + describe("InteractiveSession") { + it("should inherit the default YARN queue from LivyConf when request queue is empty") { + val testLivyConf = new LivyConf() + .set(LivyConf.REPL_JARS, "dummy.jar") + .set(LivyConf.SPARK_YARN_QUEUE, "livy-default-queue") + + val req = new CreateInteractiveRequest() + req.kind = Spark + req.queue = None + req.conf = Map(RSCConf.Entry.LIVY_JARS.key() -> "") + + val mockClient = Some(mock[RSCClient]) + + val s = InteractiveSession.create( + id = 101, + name = None, + owner = "systest", + proxyUser = None, + livyConf = testLivyConf, + accessManager = accessManager, + request = req, + sessionStore = mock[SessionStore], + ttl = None, + idleTimeout = None, + mockApp = None, + mockClient = mockClient + ) + + s.queue shouldBe Some("livy-default-queue") + } + + it("should prioritize user-specified request queue over LivyConf global configuration") { + val testLivyConf = new LivyConf() + .set(LivyConf.REPL_JARS, "dummy.jar") + .set(LivyConf.SPARK_YARN_QUEUE, "livy-default-queue") + + val req = new CreateInteractiveRequest() + req.kind = Spark + req.queue = Some("user-custom-queue") + req.conf = Map(RSCConf.Entry.LIVY_JARS.key() -> "") + + val mockClient = Some(mock[RSCClient]) + + val s = InteractiveSession.create( + id = 102, + name = None, + owner = "systest", + proxyUser = None, + livyConf = testLivyConf, + accessManager = accessManager, + request = req, + sessionStore = mock[SessionStore], + ttl = None, + idleTimeout = None, + mockApp = None, + mockClient = mockClient + ) + + s.queue shouldBe Some("user-custom-queue") + } + } }