#698 - Add support for extra yarn tags (#702)

* #698 - Add support for extra yarn tags

Author: jozefbakus

src/main/scala/za/co/absa/hyperdrive/trigger/api/rest/services/JobTemplateResolutionService.scala

@@ -16,7 +16,7 @@
 package za.co.absa.hyperdrive.trigger.api.rest.services
 
 import org.springframework.stereotype.Service
-import za.co.absa.hyperdrive.trigger.configuration.application.JobDefinitionConfig.{KeysToMerge, MergedValuesSeparator}
+import za.co.absa.hyperdrive.trigger.configuration.application.JobDefinitionConfig.{SparkExtraJavaOptions, SparkTags}
 import za.co.absa.hyperdrive.trigger.models._
 import za.co.absa.hyperdrive.trigger.api.rest.utils.Extensions.{SparkConfigList, SparkConfigMap}
 
@@ -124,8 +124,12 @@ class JobTemplateResolutionServiceImpl extends JobTemplateResolutionService {
     secondary ++ primary
 
   private def mergeSortedMapEntries(key: String, firstValue: String, secondValue: String): String =
-    if (KeysToMerge.contains(key)) {
-      s"$secondValue$MergedValuesSeparator$firstValue".trim
+    if (SparkExtraJavaOptions.KeysToMerge.contains(key)) {
+      s"$secondValue${SparkExtraJavaOptions.MergedValuesSeparator}$firstValue".trim
+    } else if (SparkTags.KeysToMerge.contains(key)) {
+      (
+        secondValue.split(SparkTags.MergedValuesSeparator) ++ firstValue.split(SparkTags.MergedValuesSeparator)
+      ).toSet[String].map(_.trim).mkString(SparkTags.MergedValuesSeparator)
     } else {
       firstValue
     }

src/main/scala/za/co/absa/hyperdrive/trigger/configuration/application/JobDefinitionConfig.scala

@@ -16,6 +16,13 @@
 package za.co.absa.hyperdrive.trigger.configuration.application
 
 object JobDefinitionConfig {
-  val KeysToMerge = Set("spark.executor.extraJavaOptions", "spark.driver.extraJavaOptions")
-  val MergedValuesSeparator = " "
+  object SparkExtraJavaOptions {
+    val KeysToMerge = Set("spark.executor.extraJavaOptions", "spark.driver.extraJavaOptions")
+    val MergedValuesSeparator = " "
+  }
+
+  object SparkTags {
+    val KeysToMerge = Set("spark.yarn.tags")
+    val MergedValuesSeparator = ","
+  }
 }

src/main/scala/za/co/absa/hyperdrive/trigger/scheduler/executors/spark/SparkClusterService.scala

@@ -15,10 +15,10 @@
 
 package za.co.absa.hyperdrive.trigger.scheduler.executors.spark
 
-import za.co.absa.hyperdrive.trigger.configuration.application.JobDefinitionConfig.{KeysToMerge, MergedValuesSeparator}
+import za.co.absa.hyperdrive.trigger.configuration.application.JobDefinitionConfig.{SparkExtraJavaOptions, SparkTags}
 import za.co.absa.hyperdrive.trigger.models.{JobInstance, SparkInstanceParameters}
 
-import scala.concurrent.{ExecutionContext, Future}
+import scala.concurrent.Future
 
 trait SparkClusterService {
   def submitJob(
@@ -32,10 +32,21 @@ trait SparkClusterService {
   protected def mergeAdditionalSparkConfig(
     globalConfig: Map[String, String],
     jobConfig: Map[String, String]
-  ): Map[String, String] =
-    KeysToMerge.map { key =>
+  ): Map[String, String] = {
+    val extraJavaOptionsMerge = SparkExtraJavaOptions.KeysToMerge.map { key =>
       val globalValue = globalConfig.getOrElse(key, "")
       val jobValue = jobConfig.getOrElse(key, "")
-      key -> s"$globalValue$MergedValuesSeparator$jobValue".trim
-    }.toMap
+      key -> s"$globalValue${SparkExtraJavaOptions.MergedValuesSeparator}$jobValue".trim
+    }
+    val tagsOptions = SparkTags.KeysToMerge.map { key =>
+      val globalValue = globalConfig.get(key)
+      val jobValue = jobConfig.get(key)
+      val value = (
+        globalValue.map(_.split(SparkTags.MergedValuesSeparator)).getOrElse(Array.empty[String]) ++
+          jobValue.map(_.split(SparkTags.MergedValuesSeparator)).getOrElse(Array.empty[String])
+      ).toSet[String].map(_.trim).mkString(SparkTags.MergedValuesSeparator)
+      if (value.nonEmpty) Some(key -> value) else None
+    }
+    (extraJavaOptionsMerge ++ tagsOptions.flatten).toMap
+  }
 }

src/main/scala/za/co/absa/hyperdrive/trigger/scheduler/executors/spark/SparkEmrClusterServiceImpl.scala

@@ -139,11 +139,11 @@ class SparkEmrClusterServiceImpl @Inject() (
   private def getSparkArgs(id: String, jobName: String, jobParameters: SparkInstanceParameters) = {
     val config = sparkConfig.emr
     val sparkSubmitConfs = Map("--deploy-mode" -> "cluster")
-    val confs = Map("spark.app.name" -> jobName, "spark.yarn.tags" -> id) ++
+    val confs = Map("spark.app.name" -> jobName) ++
       config.additionalConfs ++
       jobParameters.additionalSparkConfig.toKeyValueMap ++
       mergeAdditionalSparkConfig(
-        config.additionalConfs,
+        config.additionalConfs ++ Map("spark.yarn.tags" -> id),
         jobParameters.additionalSparkConfig.toKeyValueMap
       )
     val files = config.filesToDeploy ++ jobParameters.additionalFiles

src/main/scala/za/co/absa/hyperdrive/trigger/scheduler/executors/spark/SparkYarnClusterServiceImpl.scala

@@ -88,7 +88,6 @@ class SparkYarnClusterServiceImpl @Inject() (implicit
       .setAppResource(jobParameters.jobJar)
       .setAppName(jobName)
       .setConf("spark.app.name", jobName)
-      .setConf("spark.yarn.tags", id)
       .addAppArgs(jobParameters.appArguments.toSeq.map(fix_json_for_yarn): _*)
       .addSparkArg("--verbose")
     config.filesToDeploy.foreach(file => sparkLauncher.addFile(file))
@@ -97,10 +96,9 @@ class SparkYarnClusterServiceImpl @Inject() (implicit
     jobParameters.additionalFiles.foreach(sparkLauncher.addFile)
     jobParameters.additionalSparkConfig.foreach(conf => sparkLauncher.setConf(conf.key, conf.value))
     mergeAdditionalSparkConfig(
-      config.additionalConfs,
+      config.additionalConfs ++ Map("spark.yarn.tags" -> id),
       jobParameters.additionalSparkConfig.toKeyValueMap
-    )
-      .foreach(conf => sparkLauncher.setConf(conf._1, conf._2))
+    ).foreach(conf => sparkLauncher.setConf(conf._1, conf._2))
 
     sparkLauncher
   }

src/test/scala/za/co/absa/hyperdrive/trigger/api/rest/utils/JobTemplateResolutionServiceTest.scala

@@ -369,6 +369,41 @@ class JobTemplateResolutionServiceTest extends FlatSpec with Matchers {
     )
   }
 
+  it should "in additionalSparkConfig, concatenate the values and filter duplicates if the key is spark.yarn.tags" in {
+    // given
+    val userParameters = SparkDefinitionParameters(
+      jobType = JobTypes.Spark,
+      jobJar = None,
+      mainClass = None,
+      additionalSparkConfig = List(
+        AdditionalSparkConfig("spark.yarn.tags", "first,second,third")
+      )
+    )
+    val templateParameters = SparkTemplateParameters(
+      jobType = JobTypes.Spark,
+      jobJar = "jobJar",
+      mainClass = "mainClass",
+      additionalSparkConfig = List(
+        AdditionalSparkConfig("spark.yarn.tags", "third,first,fourth")
+      )
+    )
+
+    val jobTemplate = GenericSparkJobTemplate.copy(jobParameters = templateParameters)
+    val jobDefinition = createJobDefinition().copy(jobTemplateId = Some(jobTemplate.id), jobParameters = userParameters)
+    val dagDefinitionJoined = createDagDefinitionJoined(jobDefinition)
+
+    // when
+    val resolvedJobDefinitions = underTest.resolveDagDefinitionJoined(dagDefinitionJoined, Seq(jobTemplate))
+
+    // then
+    val resolvedJobDefinition = resolvedJobDefinitions.head
+    resolvedJobDefinition.jobParameters
+      .asInstanceOf[SparkInstanceParameters]
+      .additionalSparkConfig should contain theSameElementsAs List(
+      AdditionalSparkConfig("spark.yarn.tags", "third,first,fourth,second")
+    )
+  }
+
   it should "throw an error if the jobTemplate is of the different type as job definiton" in {
     // given
     val jobTemplate = GenericShellJobTemplate

src/test/scala/za/co/absa/hyperdrive/trigger/scheduler/executors/spark/SparkClusterServiceTest.scala

@@ -0,0 +1,103 @@
+/*
+ * Copyright 2018 ABSA Group Limited
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package za.co.absa.hyperdrive.trigger.scheduler.executors.spark
+
+import org.scalatest.{FlatSpec, Matchers}
+import za.co.absa.hyperdrive.trigger.models.{JobInstance, SparkInstanceParameters}
+
+import scala.concurrent.Future
+
+class SparkClusterServiceTest extends FlatSpec with Matchers with SparkClusterService {
+
+  override def submitJob(
+    jobInstance: JobInstance,
+    jobParameters: SparkInstanceParameters,
+    updateJob: JobInstance => Future[Unit]
+  ): Future[Unit] = Future.successful()
+
+  override def handleMissingYarnStatus(
+    jobInstance: JobInstance,
+    updateJob: JobInstance => Future[Unit]
+  ): Future[Unit] = Future.successful()
+
+  "SparkClusterService.mergeAdditionalSparkConfig" should "merge empty inputs" in {
+    val first = Map.empty[String, String]
+    val second = Map.empty[String, String]
+
+    val result = this.mergeAdditionalSparkConfig(first, second)
+
+    result should contain theSameElementsAs Map(
+      "spark.executor.extraJavaOptions" -> "",
+      "spark.driver.extraJavaOptions" -> ""
+    )
+  }
+
+  it should "concatenate the values if the key is extraJavaOptions" in {
+    val first = Map(
+      "spark.driver.extraJavaOptions" -> "-user.prop=userDriver",
+      "spark.executor.extraJavaOptions" -> "-user.prop=userExecutor"
+    )
+    val second = Map(
+      "spark.driver.extraJavaOptions" -> "-template.prop=templateDriver",
+      "spark.executor.extraJavaOptions" -> "-template.prop=templateExecutor"
+    )
+
+    val result = this.mergeAdditionalSparkConfig(first, second)
+
+    result should contain theSameElementsAs Map(
+      "spark.driver.extraJavaOptions" -> "-user.prop=userDriver -template.prop=templateDriver",
+      "spark.executor.extraJavaOptions" -> "-user.prop=userExecutor -template.prop=templateExecutor"
+    )
+  }
+
+  it should "concatenate the values if the key is spark.yarn.tags" in {
+    val first = Map(
+      "spark.yarn.tags" -> "first,second,third"
+    )
+    val second = Map(
+      "spark.yarn.tags" -> "third,first,fourth"
+    )
+
+    val result = this.mergeAdditionalSparkConfig(first, second)
+
+    result should contain theSameElementsAs Map(
+      "spark.yarn.tags" -> "first,second,third,fourth",
+      "spark.driver.extraJavaOptions" -> "",
+      "spark.executor.extraJavaOptions" -> ""
+    )
+  }
+
+  it should "concatenate the values if the key is extraJavaOptions or spark.yarn.tags" in {
+    val first = Map(
+      "spark.driver.extraJavaOptions" -> "-user.prop=userDriver",
+      "spark.executor.extraJavaOptions" -> "-user.prop=userExecutor",
+      "spark.yarn.tags" -> "first,second,third"
+    )
+    val second = Map(
+      "spark.driver.extraJavaOptions" -> "-template.prop=templateDriver",
+      "spark.executor.extraJavaOptions" -> "-template.prop=templateExecutor",
+      "spark.yarn.tags" -> "third,first,fourth"
+    )
+
+    val result = this.mergeAdditionalSparkConfig(first, second)
+
+    result should contain theSameElementsAs Map(
+      "spark.driver.extraJavaOptions" -> "-user.prop=userDriver -template.prop=templateDriver",
+      "spark.executor.extraJavaOptions" -> "-user.prop=userExecutor -template.prop=templateExecutor",
+      "spark.yarn.tags" -> "first,second,third,fourth"
+    )
+  }
+}