twill-yarn/src/main/java/org/apache/twill/yarn/YarnTwillController.java - twill - Git at Google

 /*
  * 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.twill.yarn;

 import com.google.common.base.Throwables;
 import com.google.common.collect.ImmutableList;
 import com.google.common.util.concurrent.ListenableFuture;
 import com.google.common.util.concurrent.Uninterruptibles;
 import org.apache.commons.lang.time.StopWatch;
 import org.apache.hadoop.yarn.api.records.ApplicationId;
 import org.apache.hadoop.yarn.api.records.FinalApplicationStatus;
 import org.apache.hadoop.yarn.api.records.YarnApplicationState;
 import org.apache.twill.api.ResourceReport;
 import org.apache.twill.api.RunId;
 import org.apache.twill.api.TwillController;
 import org.apache.twill.api.logging.LogHandler;
 import org.apache.twill.internal.AbstractTwillController;
 import org.apache.twill.internal.Constants;
 import org.apache.twill.internal.ProcessController;
 import org.apache.twill.internal.appmaster.TrackerService;
 import org.apache.twill.internal.state.SystemMessages;
 import org.apache.twill.internal.yarn.YarnApplicationReport;
 import org.apache.twill.zookeeper.NodeData;
 import org.apache.twill.zookeeper.ZKClient;
 import org.apache.zookeeper.data.Stat;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;

 import java.io.IOException;
 import java.net.URI;
 import java.net.URL;
 import java.util.concurrent.Callable;
 import java.util.concurrent.ExecutionException;
 import java.util.concurrent.TimeUnit;
 import java.util.concurrent.TimeoutException;

 /**
  * A {@link org.apache.twill.api.TwillController} that controllers application running on Hadoop YARN.
  */
 final class YarnTwillController extends AbstractTwillController implements TwillController {

   private static final Logger LOG = LoggerFactory.getLogger(YarnTwillController.class);

   private final String appName;
   private final Callable<ProcessController<YarnApplicationReport>> startUp;
   private ProcessController<YarnApplicationReport> processController;
   private ResourceReportClient resourcesClient;

   // Thread for polling yarn for application status if application got ZK session expire.
   // Only used by the instanceUpdate/Delete method, which is from serialized call from ZK callback.
   private Thread statusPollingThread;

   /**
    * Creates an instance without any {@link LogHandler}.
    */
   YarnTwillController(String appName, RunId runId, ZKClient zkClient,
                       Callable<ProcessController<YarnApplicationReport>> startUp) {
     this(appName, runId, zkClient, ImmutableList.<LogHandler>of(), startUp);
   }

   YarnTwillController(String appName, RunId runId, ZKClient zkClient, Iterable<LogHandler> logHandlers,
                       Callable<ProcessController<YarnApplicationReport>> startUp) {
     super(runId, zkClient, logHandlers);
     this.appName = appName;
     this.startUp = startUp;
   }


   /**
    * Sends a message to application to notify the secure store has be updated.
    */
   ListenableFuture<Void> secureStoreUpdated() {
     return sendMessage(SystemMessages.SECURE_STORE_UPDATED, null);
   }

   @Override
   protected void doStartUp() {
     super.doStartUp();

     // Submit and poll the status of the yarn application
     try {
       processController = startUp.call();

       YarnApplicationReport report = processController.getReport();
       ApplicationId appId = report.getApplicationId();
       LOG.debug("Application {} with id {} submitted", appName, appId);

       YarnApplicationState state = report.getYarnApplicationState();
       StopWatch stopWatch = new StopWatch();
       stopWatch.start();
       stopWatch.split();
       long maxTime = TimeUnit.MILLISECONDS.convert(Constants.APPLICATION_MAX_START_SECONDS, TimeUnit.SECONDS);

       LOG.debug("Checking yarn application status for {} {}", appName, appId);
       while (!hasRun(state) && stopWatch.getSplitTime() < maxTime) {
         report = processController.getReport();
         state = report.getYarnApplicationState();
         LOG.debug("Yarn application status for {} {}: {}", appName, appId, state);
         TimeUnit.SECONDS.sleep(1);
         stopWatch.split();
       }
       LOG.info("Yarn application {} {} is in state {}", appName, appId, state);
       if (state != YarnApplicationState.RUNNING) {
         LOG.info("Yarn application {} {} is not in running state. Shutting down controller.",
                  appName, appId, Constants.APPLICATION_MAX_START_SECONDS);
         forceShutDown();
       } else {
         try {
           URL resourceUrl = URI.create(String.format("http://%s:%d", report.getHost(), report.getRpcPort()))
                                .resolve(TrackerService.PATH).toURL();
           resourcesClient = new ResourceReportClient(resourceUrl);
         } catch (IOException e) {
           resourcesClient = null;
         }
       }
     } catch (Exception e) {
       throw Throwables.propagate(e);
     }
   }

   @Override
   protected synchronized void doShutDown() {
     if (processController == null) {
       LOG.warn("No process controller for application that is not submitted.");
       return;
     }

     // Stop polling if it is running.
     stopPollStatus();

     // Wait for the stop message being processed
     try {
       Uninterruptibles.getUninterruptibly(getStopMessageFuture(),
                                           Constants.APPLICATION_MAX_STOP_SECONDS, TimeUnit.SECONDS);
     } catch (Exception e) {
       LOG.error("Failed to wait for stop message being processed.", e);
       // Kill the application through yarn
       kill();
     }

     // Poll application status from yarn
     try {
       StopWatch stopWatch = new StopWatch();
       stopWatch.start();
       stopWatch.split();
       long maxTime = TimeUnit.MILLISECONDS.convert(Constants.APPLICATION_MAX_STOP_SECONDS, TimeUnit.SECONDS);

       YarnApplicationReport report = processController.getReport();
       FinalApplicationStatus finalStatus = report.getFinalApplicationStatus();
       ApplicationId appId = report.getApplicationId();
       while (finalStatus == FinalApplicationStatus.UNDEFINED && stopWatch.getSplitTime() < maxTime) {
         LOG.debug("Yarn application final status for {} {}: {}", appName, appId, finalStatus);
         TimeUnit.SECONDS.sleep(1);
         stopWatch.split();
         finalStatus = processController.getReport().getFinalApplicationStatus();
       }
       LOG.debug("Yarn application {} {} completed with status {}", appName, appId, finalStatus);

       // Application not finished after max stop time, kill the application
       if (finalStatus == FinalApplicationStatus.UNDEFINED) {
         kill();
       }
     } catch (Exception e) {
       LOG.warn("Exception while waiting for application report: {}", e.getMessage(), e);
       kill();
     }

     super.doShutDown();
   }

   @Override
   public void kill() {
     if (processController != null) {
       YarnApplicationReport report = processController.getReport();
       LOG.info("Killing application {} {}", appName, report.getApplicationId());
       processController.cancel();
     } else {
       LOG.warn("No process controller for application that is not submitted.");
     }
   }

   @Override
   protected void instanceNodeUpdated(NodeData nodeData) {
   }

   @Override
   protected void instanceNodeFailed(Throwable cause) {
     // Resort to polling from Yarn for the application status.
     if (processController == null) {
       LOG.warn("No process controller for application that is not submitted.");
       return;
     }
     YarnApplicationReport report = processController.getReport();

     // It happens if the application has ZK session expire or the node is deleted due to application termination.
     LOG.info("Failed to access application {} {} live node in ZK, resort to polling. Failure reason: {}",
              appName, report.getApplicationId(), cause == null ? "Unknown" : cause.getMessage());

     startPollStatus(report.getApplicationId());
   }

   private synchronized void startPollStatus(ApplicationId appId) {
     if (statusPollingThread == null) {
       statusPollingThread = new Thread(createStatusPollingRunnable(),
                                        String.format("%s-%s-yarn-poller", appName, appId));
       statusPollingThread.setDaemon(true);
       statusPollingThread.start();
     }
   }

   private synchronized void stopPollStatus() {
     if (statusPollingThread != null) {
       statusPollingThread.interrupt();
       statusPollingThread = null;
     }
   }

   private Runnable createStatusPollingRunnable() {
     return new Runnable() {

       @Override
       public void run() {
         YarnApplicationReport report = processController.getReport();
         ApplicationId appId = report.getApplicationId();
         boolean shutdown = false;
         boolean watchInstanceNode = false;

         try {
           LOG.debug("Polling status from Yarn for {} {}.", appName, appId);
           while (!Thread.currentThread().isInterrupted()) {
             if (report.getFinalApplicationStatus() != FinalApplicationStatus.UNDEFINED) {
               shutdown = true;
               break;
             }
             // Make a sync exists call to instance node and re-watch if the node exists
             try {
               // The timeout is arbitrary, as it's just for avoiding block forever
               Stat stat = zkClient.exists(getInstancePath()).get(5, TimeUnit.SECONDS);
               if (stat != null) {
                 watchInstanceNode = true;
                 break;
               }
             } catch (ExecutionException e) {
               // Ignore the exception, as any exception won't affect the status polling.
               LOG.debug("Failed in exists call on ZK path {}.", getInstancePath(), e);
             } catch (TimeoutException e) {
               LOG.debug("Timeout in exists call on ZK path {}.", getInstancePath(), e);
             }

             TimeUnit.SECONDS.sleep(1);
             report = processController.getReport();
           }
         } catch (InterruptedException e) {
           // OK to ignore.
           LOG.debug("Status polling thread interrupted for application {} {}", appName, appId);
         }

         LOG.debug("Stop polling status from Yarn for {} {}.", appName, appId);

         if (shutdown) {
           LOG.info("Yarn application {} {} completed. Shutting down controller.", appName, appId);
           forceShutDown();
         } else if (watchInstanceNode) {
           LOG.info("Rewatch instance node for {} {} at {}", appName, appId, getInstancePath());
           synchronized (YarnTwillController.this) {
             statusPollingThread = null;
             watchInstanceNode();
           }
         }
       }
     };
   }

   private boolean hasRun(YarnApplicationState state) {
     switch (state) {
       case RUNNING:
       case FINISHED:
       case FAILED:
       case KILLED:
         return true;
     }
     return false;
   }

   @Override
   public ResourceReport getResourceReport() {
     // in case the user calls this before starting, return null
     return (resourcesClient == null) ? null : resourcesClient.get();
   }
 }
	/*
	* 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.twill.yarn;

	import com.google.common.base.Throwables;
	import com.google.common.collect.ImmutableList;
	import com.google.common.util.concurrent.ListenableFuture;
	import com.google.common.util.concurrent.Uninterruptibles;
	import org.apache.commons.lang.time.StopWatch;
	import org.apache.hadoop.yarn.api.records.ApplicationId;
	import org.apache.hadoop.yarn.api.records.FinalApplicationStatus;
	import org.apache.hadoop.yarn.api.records.YarnApplicationState;
	import org.apache.twill.api.ResourceReport;
	import org.apache.twill.api.RunId;
	import org.apache.twill.api.TwillController;
	import org.apache.twill.api.logging.LogHandler;
	import org.apache.twill.internal.AbstractTwillController;
	import org.apache.twill.internal.Constants;
	import org.apache.twill.internal.ProcessController;
	import org.apache.twill.internal.appmaster.TrackerService;
	import org.apache.twill.internal.state.SystemMessages;
	import org.apache.twill.internal.yarn.YarnApplicationReport;
	import org.apache.twill.zookeeper.NodeData;
	import org.apache.twill.zookeeper.ZKClient;
	import org.apache.zookeeper.data.Stat;
	import org.slf4j.Logger;
	import org.slf4j.LoggerFactory;

	import java.io.IOException;
	import java.net.URI;
	import java.net.URL;
	import java.util.concurrent.Callable;
	import java.util.concurrent.ExecutionException;
	import java.util.concurrent.TimeUnit;
	import java.util.concurrent.TimeoutException;

	/**
	* A {@link org.apache.twill.api.TwillController} that controllers application running on Hadoop YARN.
	*/
	final class YarnTwillController extends AbstractTwillController implements TwillController {

	private static final Logger LOG = LoggerFactory.getLogger(YarnTwillController.class);

	private final String appName;
	private final Callable<ProcessController<YarnApplicationReport>> startUp;
	private ProcessController<YarnApplicationReport> processController;
	private ResourceReportClient resourcesClient;

	// Thread for polling yarn for application status if application got ZK session expire.
	// Only used by the instanceUpdate/Delete method, which is from serialized call from ZK callback.
	private Thread statusPollingThread;

	/**
	* Creates an instance without any {@link LogHandler}.
	*/
	YarnTwillController(String appName, RunId runId, ZKClient zkClient,
	Callable<ProcessController<YarnApplicationReport>> startUp) {
	this(appName, runId, zkClient, ImmutableList.<LogHandler>of(), startUp);
	}

	YarnTwillController(String appName, RunId runId, ZKClient zkClient, Iterable<LogHandler> logHandlers,
	Callable<ProcessController<YarnApplicationReport>> startUp) {
	super(runId, zkClient, logHandlers);
	this.appName = appName;
	this.startUp = startUp;
	}


	/**
	* Sends a message to application to notify the secure store has be updated.
	*/
	ListenableFuture<Void> secureStoreUpdated() {
	return sendMessage(SystemMessages.SECURE_STORE_UPDATED, null);
	}

	@Override
	protected void doStartUp() {
	super.doStartUp();

	// Submit and poll the status of the yarn application
	try {
	processController = startUp.call();

	YarnApplicationReport report = processController.getReport();
	ApplicationId appId = report.getApplicationId();
	LOG.debug("Application {} with id {} submitted", appName, appId);

	YarnApplicationState state = report.getYarnApplicationState();
	StopWatch stopWatch = new StopWatch();
	stopWatch.start();
	stopWatch.split();
	long maxTime = TimeUnit.MILLISECONDS.convert(Constants.APPLICATION_MAX_START_SECONDS, TimeUnit.SECONDS);

	LOG.debug("Checking yarn application status for {} {}", appName, appId);
	while (!hasRun(state) && stopWatch.getSplitTime() < maxTime) {
	report = processController.getReport();
	state = report.getYarnApplicationState();
	LOG.debug("Yarn application status for {} {}: {}", appName, appId, state);
	TimeUnit.SECONDS.sleep(1);
	stopWatch.split();
	}
	LOG.info("Yarn application {} {} is in state {}", appName, appId, state);
	if (state != YarnApplicationState.RUNNING) {
	LOG.info("Yarn application {} {} is not in running state. Shutting down controller.",
	appName, appId, Constants.APPLICATION_MAX_START_SECONDS);
	forceShutDown();
	} else {
	try {
	URL resourceUrl = URI.create(String.format("http://%s:%d", report.getHost(), report.getRpcPort()))
	.resolve(TrackerService.PATH).toURL();
	resourcesClient = new ResourceReportClient(resourceUrl);
	} catch (IOException e) {
	resourcesClient = null;
	}
	}
	} catch (Exception e) {
	throw Throwables.propagate(e);
	}
	}

	@Override
	protected synchronized void doShutDown() {
	if (processController == null) {
	LOG.warn("No process controller for application that is not submitted.");
	return;
	}

	// Stop polling if it is running.
	stopPollStatus();

	// Wait for the stop message being processed
	try {
	Uninterruptibles.getUninterruptibly(getStopMessageFuture(),
	Constants.APPLICATION_MAX_STOP_SECONDS, TimeUnit.SECONDS);
	} catch (Exception e) {
	LOG.error("Failed to wait for stop message being processed.", e);
	// Kill the application through yarn
	kill();
	}

	// Poll application status from yarn
	try {
	StopWatch stopWatch = new StopWatch();
	stopWatch.start();
	stopWatch.split();
	long maxTime = TimeUnit.MILLISECONDS.convert(Constants.APPLICATION_MAX_STOP_SECONDS, TimeUnit.SECONDS);

	YarnApplicationReport report = processController.getReport();
	FinalApplicationStatus finalStatus = report.getFinalApplicationStatus();
	ApplicationId appId = report.getApplicationId();
	while (finalStatus == FinalApplicationStatus.UNDEFINED && stopWatch.getSplitTime() < maxTime) {
	LOG.debug("Yarn application final status for {} {}: {}", appName, appId, finalStatus);
	TimeUnit.SECONDS.sleep(1);
	stopWatch.split();
	finalStatus = processController.getReport().getFinalApplicationStatus();
	}
	LOG.debug("Yarn application {} {} completed with status {}", appName, appId, finalStatus);

	// Application not finished after max stop time, kill the application
	if (finalStatus == FinalApplicationStatus.UNDEFINED) {
	kill();
	}
	} catch (Exception e) {
	LOG.warn("Exception while waiting for application report: {}", e.getMessage(), e);
	kill();
	}

	super.doShutDown();
	}

	@Override
	public void kill() {
	if (processController != null) {
	YarnApplicationReport report = processController.getReport();
	LOG.info("Killing application {} {}", appName, report.getApplicationId());
	processController.cancel();
	} else {
	LOG.warn("No process controller for application that is not submitted.");
	}
	}

	@Override
	protected void instanceNodeUpdated(NodeData nodeData) {
	}

	@Override
	protected void instanceNodeFailed(Throwable cause) {
	// Resort to polling from Yarn for the application status.
	if (processController == null) {
	LOG.warn("No process controller for application that is not submitted.");
	return;
	}
	YarnApplicationReport report = processController.getReport();

	// It happens if the application has ZK session expire or the node is deleted due to application termination.
	LOG.info("Failed to access application {} {} live node in ZK, resort to polling. Failure reason: {}",
	appName, report.getApplicationId(), cause == null ? "Unknown" : cause.getMessage());

	startPollStatus(report.getApplicationId());
	}

	private synchronized void startPollStatus(ApplicationId appId) {
	if (statusPollingThread == null) {
	statusPollingThread = new Thread(createStatusPollingRunnable(),
	String.format("%s-%s-yarn-poller", appName, appId));
	statusPollingThread.setDaemon(true);
	statusPollingThread.start();
	}
	}

	private synchronized void stopPollStatus() {
	if (statusPollingThread != null) {
	statusPollingThread.interrupt();
	statusPollingThread = null;
	}
	}

	private Runnable createStatusPollingRunnable() {
	return new Runnable() {

	@Override
	public void run() {
	YarnApplicationReport report = processController.getReport();
	ApplicationId appId = report.getApplicationId();
	boolean shutdown = false;
	boolean watchInstanceNode = false;

	try {
	LOG.debug("Polling status from Yarn for {} {}.", appName, appId);
	while (!Thread.currentThread().isInterrupted()) {
	if (report.getFinalApplicationStatus() != FinalApplicationStatus.UNDEFINED) {
	shutdown = true;
	break;
	}
	// Make a sync exists call to instance node and re-watch if the node exists
	try {
	// The timeout is arbitrary, as it's just for avoiding block forever
	Stat stat = zkClient.exists(getInstancePath()).get(5, TimeUnit.SECONDS);
	if (stat != null) {
	watchInstanceNode = true;
	break;
	}
	} catch (ExecutionException e) {
	// Ignore the exception, as any exception won't affect the status polling.
	LOG.debug("Failed in exists call on ZK path {}.", getInstancePath(), e);
	} catch (TimeoutException e) {
	LOG.debug("Timeout in exists call on ZK path {}.", getInstancePath(), e);
	}

	TimeUnit.SECONDS.sleep(1);
	report = processController.getReport();
	}
	} catch (InterruptedException e) {
	// OK to ignore.
	LOG.debug("Status polling thread interrupted for application {} {}", appName, appId);
	}

	LOG.debug("Stop polling status from Yarn for {} {}.", appName, appId);

	if (shutdown) {
	LOG.info("Yarn application {} {} completed. Shutting down controller.", appName, appId);
	forceShutDown();
	} else if (watchInstanceNode) {
	LOG.info("Rewatch instance node for {} {} at {}", appName, appId, getInstancePath());
	synchronized (YarnTwillController.this) {
	statusPollingThread = null;
	watchInstanceNode();
	}
	}
	}
	};
	}

	private boolean hasRun(YarnApplicationState state) {
	switch (state) {
	case RUNNING:
	case FINISHED:
	case FAILED:
	case KILLED:
	return true;
	}
	return false;
	}

	@Override
	public ResourceReport getResourceReport() {
	// in case the user calls this before starting, return null
	return (resourcesClient == null) ? null : resourcesClient.get();
	}
	}