| /* |
| * 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.pig.backend.hadoop.executionengine.tez; |
| |
| import java.io.IOException; |
| import java.text.SimpleDateFormat; |
| import java.util.ArrayList; |
| import java.util.Calendar; |
| import java.util.Iterator; |
| import java.util.List; |
| import java.util.Map; |
| import java.util.concurrent.locks.ReentrantReadWriteLock; |
| |
| import org.apache.commons.logging.Log; |
| import org.apache.commons.logging.LogFactory; |
| import org.apache.hadoop.conf.Configuration; |
| import org.apache.hadoop.security.Credentials; |
| import org.apache.hadoop.yarn.api.records.LocalResource; |
| import org.apache.pig.PigConfiguration; |
| import org.apache.pig.backend.hadoop.executionengine.tez.TezJob.TezJobConfig; |
| import org.apache.pig.backend.hadoop.executionengine.tez.util.MRToTezHelper; |
| import org.apache.pig.impl.PigContext; |
| import org.apache.pig.impl.PigImplConstants; |
| import org.apache.pig.impl.util.Utils; |
| import org.apache.pig.tools.pigstats.tez.TezScriptState; |
| import org.apache.tez.client.TezAppMasterStatus; |
| import org.apache.tez.client.TezClient; |
| import org.apache.tez.dag.api.TezConfiguration; |
| import org.apache.tez.dag.api.TezException; |
| |
| import com.google.common.annotations.VisibleForTesting; |
| |
| public class TezSessionManager { |
| |
| private static final Log log = LogFactory.getLog(TezSessionManager.class); |
| |
| static { |
| Utils.addShutdownHookWithPriority(new Runnable() { |
| |
| @Override |
| public void run() { |
| TezSessionManager.shutdown(); |
| } |
| }, PigImplConstants.SHUTDOWN_HOOK_JOB_KILL_PRIORITY); |
| } |
| |
| private static ReentrantReadWriteLock sessionPoolLock = new ReentrantReadWriteLock(); |
| private static boolean shutdown = false; |
| |
| private TezSessionManager() { |
| } |
| |
| private static class SessionInfo { |
| |
| public SessionInfo(TezClient session, TezConfiguration config, Map<String, LocalResource> resources) { |
| this.session = session; |
| this.config = config; |
| this.resources = resources; |
| } |
| |
| public TezConfiguration getConfig() { |
| return config; |
| } |
| public Map<String, LocalResource> getResources() { |
| return resources; |
| } |
| public TezClient getTezSession() { |
| return session; |
| } |
| public void setInUse(boolean inUse) { |
| this.inUse = inUse; |
| } |
| private TezClient session; |
| private Map<String, LocalResource> resources; |
| private TezConfiguration config; |
| private boolean inUse = false; |
| } |
| |
| private static List<SessionInfo> sessionPool = new ArrayList<SessionInfo>(); |
| |
| private static SessionInfo createSession(TezConfiguration amConf, |
| Map<String, LocalResource> requestedAMResources, Credentials creds, |
| TezJobConfig tezJobConf) throws TezException, IOException, |
| InterruptedException { |
| MRToTezHelper.translateMRSettingsForTezAM(amConf); |
| TezScriptState ss = TezScriptState.get(); |
| ss.addDAGSettingsToConf(amConf); |
| if (amConf.getBoolean(PigConfiguration.PIG_TEZ_CONFIGURE_AM_MEMORY, true)) { |
| adjustAMConfig(amConf, tezJobConf); |
| } |
| String jobName = amConf.get(PigContext.JOB_NAME, "pig"); |
| TezClient tezClient = TezClient.create(jobName, amConf, true, requestedAMResources, creds); |
| try { |
| tezClient.start(); |
| TezAppMasterStatus appMasterStatus = tezClient.getAppMasterStatus(); |
| if (appMasterStatus.equals(TezAppMasterStatus.SHUTDOWN)) { |
| throw new RuntimeException("TezSession has already shutdown"); |
| } |
| tezClient.waitTillReady(); |
| } catch (Throwable e) { |
| log.error("Exception while waiting for Tez client to be ready", e); |
| tezClient.stop(); |
| throw new RuntimeException(e); |
| } |
| return new SessionInfo(tezClient, amConf, requestedAMResources); |
| } |
| |
| private static void adjustAMConfig(TezConfiguration amConf, TezJobConfig tezJobConf) { |
| String amLaunchOpts = amConf.get( |
| TezConfiguration.TEZ_AM_LAUNCH_CMD_OPTS, |
| TezConfiguration.TEZ_AM_LAUNCH_CMD_OPTS_DEFAULT); |
| int configuredAMMaxHeap = Utils.extractHeapSizeInMB(amLaunchOpts); |
| int configuredAMResourceMB = amConf.getInt( |
| TezConfiguration.TEZ_AM_RESOURCE_MEMORY_MB, |
| TezConfiguration.TEZ_AM_RESOURCE_MEMORY_MB_DEFAULT); |
| |
| if (tezJobConf.getEstimatedTotalParallelism() > 0) { |
| |
| // Need more room for native memory/virtual address space |
| // when close to 4G due to 32-bit jvm 4G limit |
| int maxAMHeap = Utils.is64bitJVM() ? 3584 : 3200; |
| int maxAMResourceMB = 4096; |
| int requiredAMResourceMB = maxAMResourceMB; |
| int requiredAMMaxHeap = maxAMHeap; |
| |
| // Rough estimation. For 5K tasks 1G Xmx and 1.5G resource.mb |
| // Increment container size by 512 mb for every additional 5K tasks. |
| // 30000 and above - 3200Xmx, 4096 (896 native memory) |
| // 25000 and above - 3072Xmx, 3584 |
| // 20000 and above - 2560Xmx, 3072 |
| // 15000 and above - 2048Xmx, 2560 |
| // 10000 and above - 1536Xmx, 2048 |
| // 5000 and above - 1024Xmx, 1536 (512 native memory) |
| for (int taskCount = 30000; taskCount >= 5000; taskCount-=5000) { |
| if (tezJobConf.getEstimatedTotalParallelism() >= taskCount) { |
| break; |
| } |
| requiredAMResourceMB = requiredAMResourceMB - 512; |
| requiredAMMaxHeap = requiredAMResourceMB - 512; |
| } |
| |
| if (tezJobConf.getTotalVertices() > 30) { |
| //Add 512 mb per 30 vertices |
| int additionaMem = 512 * (tezJobConf.getTotalVertices() / 30); |
| requiredAMResourceMB = requiredAMResourceMB + additionaMem; |
| requiredAMMaxHeap = requiredAMResourceMB - 512; |
| } |
| |
| if (tezJobConf.getMaxOutputsinSingleVertex() > 10) { |
| //Add 256 mb per 5 outputs if a vertex has more than 10 outputs |
| int additionaMem = 256 * (tezJobConf.getMaxOutputsinSingleVertex() / 5); |
| requiredAMResourceMB = requiredAMResourceMB + additionaMem; |
| requiredAMMaxHeap = requiredAMResourceMB - 512; |
| } |
| |
| requiredAMResourceMB = Math.min(maxAMResourceMB, requiredAMResourceMB); |
| requiredAMMaxHeap = Math.min(maxAMHeap, requiredAMMaxHeap); |
| |
| if (requiredAMResourceMB > -1 && configuredAMResourceMB < requiredAMResourceMB) { |
| amConf.setInt(TezConfiguration.TEZ_AM_RESOURCE_MEMORY_MB, requiredAMResourceMB); |
| log.info("Increasing " |
| + TezConfiguration.TEZ_AM_RESOURCE_MEMORY_MB + " from " |
| + configuredAMResourceMB + " to " |
| + requiredAMResourceMB |
| + " as total estimated tasks = " + tezJobConf.getEstimatedTotalParallelism() |
| + ", total vertices = " + tezJobConf.getTotalVertices() |
| + ", max outputs = " + tezJobConf.getMaxOutputsinSingleVertex()); |
| |
| if (requiredAMMaxHeap > -1 && configuredAMMaxHeap < requiredAMMaxHeap) { |
| amConf.set(TezConfiguration.TEZ_AM_LAUNCH_CMD_OPTS, |
| amLaunchOpts + " -Xmx" + requiredAMMaxHeap + "M"); |
| log.info("Increasing Tez AM Heap Size from " |
| + configuredAMMaxHeap + "M to " |
| + requiredAMMaxHeap |
| + "M as total estimated tasks = " + tezJobConf.getEstimatedTotalParallelism() |
| + ", total vertices = " + tezJobConf.getTotalVertices() |
| + ", max outputs = " + tezJobConf.getMaxOutputsinSingleVertex()); |
| log.info("Value of " + TezConfiguration.TEZ_AM_LAUNCH_CMD_OPTS + " is now " |
| + amConf.get(TezConfiguration.TEZ_AM_LAUNCH_CMD_OPTS)); |
| } |
| } |
| } |
| } |
| |
| private static boolean validateSessionResources(SessionInfo currentSession, |
| Map<String, LocalResource> requestedAMResources) |
| throws TezException, IOException { |
| for (Map.Entry<String, LocalResource> entry : requestedAMResources.entrySet()) { |
| if (!currentSession.resources.entrySet().contains(entry)) { |
| return false; |
| } |
| } |
| return true; |
| } |
| |
| private static boolean validateSessionConfig(SessionInfo currentSession, |
| Configuration newSessionConfig) |
| throws TezException, IOException { |
| // If DAG recovery is disabled for one and enabled for another, do not reuse |
| if (currentSession.getConfig().getBoolean( |
| TezConfiguration.DAG_RECOVERY_ENABLED, |
| TezConfiguration.DAG_RECOVERY_ENABLED_DEFAULT) |
| != newSessionConfig.getBoolean( |
| TezConfiguration.DAG_RECOVERY_ENABLED, |
| TezConfiguration.DAG_RECOVERY_ENABLED_DEFAULT)) { |
| return false; |
| } |
| return true; |
| } |
| |
| static TezClient getClient(TezConfiguration conf, Map<String, LocalResource> requestedAMResources, |
| Credentials creds, TezJobConfig tezJobConf) throws TezException, IOException, InterruptedException { |
| List<SessionInfo> sessionsToRemove = new ArrayList<SessionInfo>(); |
| SessionInfo newSession = null; |
| sessionPoolLock.readLock().lock(); |
| try { |
| if (shutdown == true) { |
| throw new IOException("TezSessionManager is shut down"); |
| } |
| |
| for (SessionInfo sessionInfo : sessionPool) { |
| synchronized (sessionInfo) { |
| TezAppMasterStatus appMasterStatus = sessionInfo.session |
| .getAppMasterStatus(); |
| if (appMasterStatus.equals(TezAppMasterStatus.SHUTDOWN)) { |
| sessionsToRemove.add(sessionInfo); |
| } else if (!sessionInfo.inUse |
| && appMasterStatus.equals(TezAppMasterStatus.READY) |
| && validateSessionResources(sessionInfo,requestedAMResources) |
| && validateSessionConfig(sessionInfo, conf)) { |
| sessionInfo.inUse = true; |
| return sessionInfo.session; |
| } |
| } |
| } |
| } finally { |
| sessionPoolLock.readLock().unlock(); |
| } |
| // We cannot find available AM, create new one |
| // Create session outside of locks so that getClient/freeSession is not |
| // blocked for parallel embedded pig runs |
| newSession = createSession(conf, requestedAMResources, creds, tezJobConf); |
| newSession.inUse = true; |
| sessionPoolLock.writeLock().lock(); |
| try { |
| if (shutdown == true) { |
| log.info("Shutting down Tez session " + newSession.session); |
| newSession.session.stop(); |
| throw new IOException("TezSessionManager is shut down"); |
| } |
| sessionPool.add(newSession); |
| for (SessionInfo sessionToRemove : sessionsToRemove) { |
| sessionPool.remove(sessionToRemove); |
| } |
| return newSession.session; |
| } finally { |
| sessionPoolLock.writeLock().unlock(); |
| } |
| } |
| |
| static void freeSession(TezClient session) { |
| sessionPoolLock.readLock().lock(); |
| try { |
| for (SessionInfo sessionInfo : sessionPool) { |
| synchronized (sessionInfo) { |
| if (sessionInfo.session == session) { |
| sessionInfo.inUse = false; |
| break; |
| } |
| } |
| } |
| } finally { |
| sessionPoolLock.readLock().unlock(); |
| } |
| } |
| |
| static void stopSession(TezClient session) throws TezException, IOException { |
| Iterator<SessionInfo> iter = sessionPool.iterator(); |
| SessionInfo sessionToRemove = null; |
| sessionPoolLock.readLock().lock(); |
| try { |
| while (iter.hasNext()) { |
| SessionInfo sessionInfo = iter.next(); |
| synchronized (sessionInfo) { |
| if (sessionInfo.session == session) { |
| log.info("Stopping Tez session " + session); |
| String timeStamp = new SimpleDateFormat("yyyy-MM-dd HH:mm:ss") |
| .format(Calendar.getInstance().getTime()); |
| System.err.println(timeStamp + " Shutting down Tez session " |
| + ", sessionName=" + session.getClientName() |
| + ", applicationId=" + session.getAppMasterApplicationId()); |
| session.stop(); |
| sessionToRemove = sessionInfo; |
| break; |
| } |
| } |
| } |
| } finally { |
| sessionPoolLock.readLock().unlock(); |
| } |
| if (sessionToRemove != null) { |
| sessionPoolLock.writeLock().lock(); |
| try { |
| sessionPool.remove(sessionToRemove); |
| } finally { |
| sessionPoolLock.writeLock().unlock(); |
| } |
| } |
| } |
| |
| @VisibleForTesting |
| public static void shutdown() { |
| try { |
| sessionPoolLock.writeLock().lock(); |
| shutdown = true; |
| for (SessionInfo sessionInfo : sessionPool) { |
| synchronized (sessionInfo) { |
| TezClient session = sessionInfo.session; |
| try { |
| String timeStamp = new SimpleDateFormat( |
| "yyyy-MM-dd HH:mm:ss").format(Calendar.getInstance().getTime()); |
| if (session.getAppMasterStatus().equals( |
| TezAppMasterStatus.SHUTDOWN)) { |
| log.info("Tez session is already shutdown " |
| + session); |
| System.err.println(timeStamp |
| + " Tez session is already shutdown " + session |
| + ", sessionName=" + session.getClientName() |
| + ", applicationId=" + session.getAppMasterApplicationId()); |
| continue; |
| } |
| log.info("Shutting down Tez session " + session); |
| // Since hadoop calls org.apache.log4j.LogManager.shutdown(); |
| // the log.info message is not displayed with shutdown hook in Oozie |
| System.err.println(timeStamp + " Shutting down Tez session " |
| + ", sessionName=" + session.getClientName() |
| + ", applicationId=" + session.getAppMasterApplicationId()); |
| session.stop(); |
| } catch (Exception e) { |
| log.error("Error shutting down Tez session " |
| + session, e); |
| } |
| } |
| } |
| sessionPool.clear(); |
| } finally { |
| sessionPoolLock.writeLock().unlock(); |
| } |
| } |
| } |
| |
