org.apache.hadoop.yarn.client.api.YarnClient Scala Examples

The following examples show how to use org.apache.hadoop.yarn.client.api.YarnClient. You can vote up the ones you like or vote down the ones you don't like, and go to the original project or source file by following the links above each example.
Example 1
Source File: KyuubiHadoopUtil.scala    From kyuubi   with Apache License 2.0 5 votes vote down vote up 
package yaooqinn.kyuubi.utils

import java.security.PrivilegedExceptionAction
import java.util.EnumSet

import scala.collection.JavaConverters._
import scala.util.control.NonFatal

import org.apache.hadoop.security.UserGroupInformation
import org.apache.hadoop.yarn.api.records.YarnApplicationState._
import org.apache.hadoop.yarn.client.api.YarnClient
import org.apache.hadoop.yarn.conf.YarnConfiguration
import org.apache.kyuubi.Logging
import org.apache.spark.KyuubiSparkUtil

private[kyuubi] object KyuubiHadoopUtil extends Logging {

  private def createYarnClient: YarnClient = {
    val c = YarnClient.createYarnClient()
    c.init(new YarnConfiguration())
    c.start()
    c
  }

  def killYarnAppByName(appName: String): Unit = {
    val client = createYarnClient
    client.getApplications(Set("SPARK").asJava, EnumSet.of(ACCEPTED, SUBMITTED, RUNNING)).asScala
      .filter(applicationReport => applicationReport.getName.equals(appName))
      .foreach { applicationReport =>
        client.killApplication(applicationReport.getApplicationId)
      }
  }

  def doAs[T](user: UserGroupInformation)(f: => T): T = {
    try {
      user.doAs(new PrivilegedExceptionAction[T] {
        override def run(): T = f
      })
    } catch {
      case NonFatal(e) => throw KyuubiSparkUtil.findCause(e)
    }
  }

  def doAsAndLogNonFatal(user: UserGroupInformation)(f: => Unit): Unit = {
    try {
      doAs(user)(f)
    } catch {
      case NonFatal(e) => error(s"Failed to operate as ${user.getShortUserName}", e)
    }
  }

  
  def doAsRealUser[T](f: => T): T = {
    val currentUser = UserGroupInformation.getCurrentUser
    val realUser = Option(currentUser.getRealUser).getOrElse(currentUser)
    doAs(realUser)(f)
  }
}