/*
* Licensed to the Apache Software Foundation (ASF) under one or more
* contributor license agreements. See the NOTICE file distributed with
* this work for additional information regarding copyright ownership.
* The ASF licenses this file to You 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 org.apache.hive.spark.client;
import com.google.common.base.Throwables;
import com.google.common.io.Files;
import io.netty.channel.ChannelHandlerContext;
import io.netty.channel.nio.NioEventLoopGroup;
import java.io.File;
import java.io.IOException;
import java.io.Serializable;
import java.util.Iterator;
import java.util.List;
import java.util.Map;
import java.util.Set;
import java.util.concurrent.Callable;
import java.util.concurrent.CopyOnWriteArrayList;
import java.util.concurrent.ExecutionException;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors;
import java.util.concurrent.Future;
import java.util.concurrent.atomic.AtomicInteger;
import org.apache.commons.io.FileUtils;
import org.apache.hadoop.hive.common.classification.InterfaceAudience;
import org.apache.hive.spark.client.metrics.Metrics;
import org.apache.hive.spark.client.rpc.Rpc;
import org.apache.hive.spark.client.rpc.RpcConfiguration;
import org.apache.hive.spark.counter.SparkCounters;
import org.apache.spark.SparkConf;
import org.apache.spark.SparkJobInfo;
import org.apache.spark.api.java.JavaFutureAction;
import org.apache.spark.api.java.JavaSparkContext;
import org.apache.spark.scheduler.SparkListener;
import org.apache.spark.scheduler.SparkListenerJobEnd;
import org.apache.spark.scheduler.SparkListenerJobStart;
import org.apache.spark.scheduler.SparkListenerTaskEnd;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import scala.Tuple2;
import com.google.common.base.Joiner;
import com.google.common.base.Preconditions;
import com.google.common.collect.Lists;
import com.google.common.collect.Maps;
import com.google.common.util.concurrent.ThreadFactoryBuilder;
/**
* Driver code for the Spark client library.
*/
@InterfaceAudience.Private
public class RemoteDriver {
private static final Logger LOG = LoggerFactory.getLogger(RemoteDriver.class);
private final Map<String, JobWrapper<?>> activeJobs;
private final Object jcLock;
private final Object shutdownLock;
private final ExecutorService executor;
private final NioEventLoopGroup egroup;
private final Rpc clientRpc;
private final DriverProtocol protocol;
// a local temp dir specific to this driver
private final File localTmpDir;
// Used to queue up requests while the SparkContext is being created.
private final List<JobWrapper<?>> jobQueue = Lists.newLinkedList();
// jc is effectively final, but it has to be volatile since it's accessed by different
// threads while the constructor is running.
private volatile JobContextImpl jc;
private volatile boolean running;
private RemoteDriver(String[] args) throws Exception {
this.activeJobs = Maps.newConcurrentMap();
this.jcLock = new Object();
this.shutdownLock = new Object();
localTmpDir = Files.createTempDir();
SparkConf conf = new SparkConf();
String serverAddress = null;
int serverPort = -1;
for (int idx = 0; idx < args.length; idx += 2) {
String key = args[idx];
if (key.equals("--remote-host")) {
serverAddress = getArg(args, idx);
} else if (key.equals("--remote-port")) {
serverPort = Integer.parseInt(getArg(args, idx));
} else if (key.equals("--client-id")) {
conf.set(SparkClientFactory.CONF_CLIENT_ID, getArg(args, idx));
} else if (key.equals("--secret")) {
conf.set(SparkClientFactory.CONF_KEY_SECRET, getArg(args, idx));
} else if (key.equals("--conf")) {
String[] val = getArg(args, idx).split("[=]", 2);
conf.set(val[0], val[1]);
} else {
throw new IllegalArgumentException("Invalid command line: "
+ Joiner.on(" ").join(args));
}
}
executor = Executors.newCachedThreadPool();
LOG.info("Connecting to: {}:{}", serverAddress, serverPort);
Map<String, String> mapConf = Maps.newHashMap();
for (Tuple2<String, String> e : conf.getAll()) {
mapConf.put(e._1(), e._2());
LOG.debug("Remote Driver configured with: " + e._1() + "=" + e._2());
}
String clientId = mapConf.get(SparkClientFactory.CONF_CLIENT_ID);
Preconditions.checkArgument(clientId != null, "No client ID provided.");
String secret = mapConf.get(SparkClientFactory.CONF_KEY_SECRET);
Preconditions.checkArgument(secret != null, "No secret provided.");
int threadCount = new RpcConfiguration(mapConf).getRpcThreadCount();
this.egroup = new NioEventLoopGroup(
threadCount,
new ThreadFactoryBuilder()
.setNameFormat("Driver-RPC-Handler-%d")
.setDaemon(true)
.build());
this.protocol = new DriverProtocol();
// The RPC library takes care of timing out this.
this.clientRpc = Rpc.createClient(mapConf, egroup, serverAddress, serverPort,
clientId, secret, protocol).get();
this.running = true;
this.clientRpc.addListener(new Rpc.Listener() {
@Override
public void rpcClosed(Rpc rpc) {
LOG.warn("Shutting down driver because RPC channel was closed.");
shutdown(null);
}
});
try {
JavaSparkContext sc = new JavaSparkContext(conf);
sc.sc().addSparkListener(new ClientListener());
synchronized (jcLock) {
jc = new JobContextImpl(sc, localTmpDir);
jcLock.notifyAll();
}
} catch (Exception e) {
LOG.error("Failed to start SparkContext: " + e, e);
shutdown(e);
synchronized (jcLock) {
jcLock.notifyAll();
}
throw e;
}
synchronized (jcLock) {
for (Iterator<JobWrapper<?>> it = jobQueue.iterator(); it.hasNext();) {
it.next().submit();
}
}
}
private void run() throws InterruptedException {
synchronized (shutdownLock) {
while (running) {
shutdownLock.wait();
}
}
executor.shutdownNow();
try {
FileUtils.deleteDirectory(localTmpDir);
} catch (IOException e) {
LOG.warn("Failed to delete local tmp dir: " + localTmpDir, e);
}
}
private void submit(JobWrapper<?> job) {
synchronized (jcLock) {
if (jc != null) {
job.submit();
} else {
LOG.info("SparkContext not yet up, queueing job request.");
jobQueue.add(job);
}
}
}
private synchronized void shutdown(Throwable error) {
if (running) {
if (error == null) {
LOG.info("Shutting down remote driver.");
} else {
LOG.error("Shutting down remote driver due to error: " + error, error);
}
running = false;
for (JobWrapper<?> job : activeJobs.values()) {
cancelJob(job);
}
if (error != null) {
protocol.sendError(error);
}
if (jc != null) {
jc.stop();
}
clientRpc.close();
egroup.shutdownGracefully();
synchronized (shutdownLock) {
shutdownLock.notifyAll();
}
}
}
private boolean cancelJob(JobWrapper<?> job) {
boolean cancelled = false;
for (JavaFutureAction<?> action : job.jobs) {
cancelled |= action.cancel(true);
}
return cancelled | (job.future != null && job.future.cancel(true));
}
private String getArg(String[] args, int keyIdx) {
int valIdx = keyIdx + 1;
if (args.length <= valIdx) {
throw new IllegalArgumentException("Invalid command line: "
+ Joiner.on(" ").join(args));
}
return args[valIdx];
}
private class DriverProtocol extends BaseProtocol {
void sendError(Throwable error) {
LOG.debug("Send error to Client: {}", Throwables.getStackTraceAsString(error));
clientRpc.call(new Error(error));
}
<T extends Serializable> void jobFinished(String jobId, T result,
Throwable error, SparkCounters counters) {
LOG.debug("Send job({}) result to Client.", jobId);
clientRpc.call(new JobResult(jobId, result, error, counters));
}
void jobStarted(String jobId) {
clientRpc.call(new JobStarted(jobId));
}
void jobSubmitted(String jobId, int sparkJobId) {
LOG.debug("Send job({}/{}) submitted to Client.", jobId, sparkJobId);
clientRpc.call(new JobSubmitted(jobId, sparkJobId));
}
void sendMetrics(String jobId, int sparkJobId, int stageId, long taskId, Metrics metrics) {
LOG.debug("Send task({}/{}/{}/{}) metric to Client.", jobId, sparkJobId, stageId, taskId);
clientRpc.call(new JobMetrics(jobId, sparkJobId, stageId, taskId, metrics));
}
private void handle(ChannelHandlerContext ctx, CancelJob msg) {
JobWrapper<?> job = activeJobs.get(msg.id);
if (job == null || !cancelJob(job)) {
LOG.info("Requested to cancel an already finished job.");
}
}
private void handle(ChannelHandlerContext ctx, EndSession msg) {
LOG.debug("Shutting down due to EndSession request.");
shutdown(null);
}
private void handle(ChannelHandlerContext ctx, JobRequest msg) {
LOG.info("Received job request {}", msg.id);
JobWrapper<?> wrapper = new JobWrapper<Serializable>(msg);
activeJobs.put(msg.id, wrapper);
submit(wrapper);
}
private Object handle(ChannelHandlerContext ctx, SyncJobRequest msg) throws Exception {
// In case the job context is not up yet, let's wait, since this is supposed to be a
// "synchronous" RPC.
if (jc == null) {
synchronized (jcLock) {
while (jc == null) {
jcLock.wait();
if (!running) {
throw new IllegalStateException("Remote context is shutting down.");
}
}
}
}
jc.setMonitorCb(new MonitorCallback() {
@Override
public void call(JavaFutureAction<?> future,
SparkCounters sparkCounters, Set<Integer> cachedRDDIds) {
throw new IllegalStateException(
"JobContext.monitor() is not available for synchronous jobs.");
}
});
try {
return msg.job.call(jc);
} finally {
jc.setMonitorCb(null);
}
}
}
private class JobWrapper<T extends Serializable> implements Callable<Void> {
private final BaseProtocol.JobRequest<T> req;
private final List<JavaFutureAction<?>> jobs;
private final AtomicInteger jobEndReceived;
private int completed;
private SparkCounters sparkCounters;
private Set<Integer> cachedRDDIds;
private Integer sparkJobId;
private Future<?> future;
JobWrapper(BaseProtocol.JobRequest<T> req) {
this.req = req;
this.jobs = Lists.newArrayList();
completed = 0;
jobEndReceived = new AtomicInteger(0);
this.sparkCounters = null;
this.cachedRDDIds = null;
this.sparkJobId = null;
}
@Override
public Void call() throws Exception {
protocol.jobStarted(req.id);
try {
jc.setMonitorCb(new MonitorCallback() {
@Override
public void call(JavaFutureAction<?> future,
SparkCounters sparkCounters, Set<Integer> cachedRDDIds) {
monitorJob(future, sparkCounters, cachedRDDIds);
}
});
T result = req.job.call(jc);
// In case the job is empty, there won't be JobStart/JobEnd events. The only way
// to know if the job has finished is to check the futures here ourselves.
for (JavaFutureAction<?> future : jobs) {
future.get();
completed++;
LOG.debug("Client job {}: {} of {} Spark jobs finished.",
req.id, completed, jobs.size());
}
// If the job is not empty (but runs fast), we have to wait until all the TaskEnd/JobEnd
// events are processed. Otherwise, task metrics may get lost. See HIVE-13525.
if (sparkJobId != null) {
SparkJobInfo sparkJobInfo = jc.sc().statusTracker().getJobInfo(sparkJobId);
if (sparkJobInfo != null && sparkJobInfo.stageIds() != null &&
sparkJobInfo.stageIds().length > 0) {
synchronized (jobEndReceived) {
while (jobEndReceived.get() < jobs.size()) {
jobEndReceived.wait();
}
}
}
}
SparkCounters counters = null;
if (sparkCounters != null) {
counters = sparkCounters.snapshot();
}
protocol.jobFinished(req.id, result, null, counters);
} catch (Throwable t) {
// Catch throwables in a best-effort to report job status back to the client. It's
// re-thrown so that the executor can destroy the affected thread (or the JVM can
// die or whatever would happen if the throwable bubbled up).
LOG.info("Failed to run job " + req.id, t);
protocol.jobFinished(req.id, null, t,
sparkCounters != null ? sparkCounters.snapshot() : null);
throw new ExecutionException(t);
} finally {
jc.setMonitorCb(null);
activeJobs.remove(req.id);
releaseCache();
}
return null;
}
void submit() {
this.future = executor.submit(this);
}
void jobDone() {
synchronized (jobEndReceived) {
jobEndReceived.incrementAndGet();
jobEndReceived.notifyAll();
}
}
/**
* Release cached RDDs as soon as the job is done.
* This is different from local Spark client so as
* to save a RPC call/trip, avoid passing cached RDD
* id information around. Otherwise, we can follow
* the local Spark client way to be consistent.
*/
void releaseCache() {
if (cachedRDDIds != null) {
for (Integer cachedRDDId: cachedRDDIds) {
jc.sc().sc().unpersistRDD(cachedRDDId, false);
}
}
}
private void monitorJob(JavaFutureAction<?> job,
SparkCounters sparkCounters, Set<Integer> cachedRDDIds) {
jobs.add(job);
if (!jc.getMonitoredJobs().containsKey(req.id)) {
jc.getMonitoredJobs().put(req.id, new CopyOnWriteArrayList<JavaFutureAction<?>>());
}
jc.getMonitoredJobs().get(req.id).add(job);
this.sparkCounters = sparkCounters;
this.cachedRDDIds = cachedRDDIds;
sparkJobId = job.jobIds().get(0);
protocol.jobSubmitted(req.id, sparkJobId);
}
}
private class ClientListener extends SparkListener {
private final Map<Integer, Integer> stageToJobId = Maps.newHashMap();
@Override
public void onJobStart(SparkListenerJobStart jobStart) {
synchronized (stageToJobId) {
for (int i = 0; i < jobStart.stageIds().length(); i++) {
stageToJobId.put((Integer) jobStart.stageIds().apply(i), jobStart.jobId());
}
}
}
@Override
public void onJobEnd(SparkListenerJobEnd jobEnd) {
synchronized (stageToJobId) {
for (Iterator<Map.Entry<Integer, Integer>> it = stageToJobId.entrySet().iterator();
it.hasNext();) {
Map.Entry<Integer, Integer> e = it.next();
if (e.getValue() == jobEnd.jobId()) {
it.remove();
}
}
}
String clientId = getClientId(jobEnd.jobId());
if (clientId != null) {
activeJobs.get(clientId).jobDone();
}
}
@Override
public void onTaskEnd(SparkListenerTaskEnd taskEnd) {
if (taskEnd.reason() instanceof org.apache.spark.Success$
&& !taskEnd.taskInfo().speculative()) {
Metrics metrics = new Metrics(taskEnd.taskMetrics());
Integer jobId;
synchronized (stageToJobId) {
jobId = stageToJobId.get(taskEnd.stageId());
}
// TODO: implement implicit AsyncRDDActions conversion instead of jc.monitor()?
// TODO: how to handle stage failures?
String clientId = getClientId(jobId);
if (clientId != null) {
protocol.sendMetrics(clientId, jobId, taskEnd.stageId(),
taskEnd.taskInfo().taskId(), metrics);
}
}
}
/**
* Returns the client job ID for the given Spark job ID.
*
* This will only work for jobs monitored via JobContext#monitor(). Other jobs won't be
* matched, and this method will return `None`.
*/
private String getClientId(Integer jobId) {
for (Map.Entry<String, JobWrapper<?>> e : activeJobs.entrySet()) {
for (JavaFutureAction<?> future : e.getValue().jobs) {
if (future.jobIds().contains(jobId)) {
return e.getKey();
}
}
}
return null;
}
}
public static void main(String[] args) throws Exception {
new RemoteDriver(args).run();
}
}