/** * 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.hadoop.yarn.server.resourcemanager; import java.io.IOException; import java.nio.ByteBuffer; import java.util.ArrayList; import java.util.HashMap; import java.util.Map; import org.apache.commons.logging.Log; import org.apache.commons.logging.LogFactory; import org.apache.hadoop.yarn.api.ApplicationConstants; import org.apache.hadoop.yarn.api.ContainerManagementProtocol; import org.apache.hadoop.yarn.api.protocolrecords.AllocateResponse; import org.apache.hadoop.yarn.api.protocolrecords.GetContainerStatusesRequest; import org.apache.hadoop.yarn.api.protocolrecords.GetContainerStatusesResponse; import org.apache.hadoop.yarn.api.protocolrecords.StartContainerRequest; import org.apache.hadoop.yarn.api.protocolrecords.StartContainersRequest; import org.apache.hadoop.yarn.api.protocolrecords.StartContainersResponse; import org.apache.hadoop.yarn.api.protocolrecords.StopContainersRequest; import org.apache.hadoop.yarn.api.protocolrecords.StopContainersResponse; import org.apache.hadoop.yarn.api.records.ApplicationAttemptId; import org.apache.hadoop.yarn.api.records.ContainerId; import org.apache.hadoop.yarn.api.records.ContainerState; import org.apache.hadoop.yarn.api.records.ResourceRequest; import org.apache.hadoop.yarn.api.records.SerializedException; import org.apache.hadoop.yarn.api.records.Token; import org.apache.hadoop.yarn.conf.YarnConfiguration; import org.apache.hadoop.yarn.exceptions.YarnException; import org.apache.hadoop.yarn.ipc.RPCUtil; import org.apache.hadoop.yarn.security.ContainerTokenIdentifier; import org.apache.hadoop.yarn.server.resourcemanager.rmapp.RMApp; import org.apache.hadoop.yarn.server.resourcemanager.rmapp.attempt.RMAppAttempt; import org.apache.hadoop.yarn.server.resourcemanager.rmapp.attempt.RMAppAttemptState; import org.apache.hadoop.yarn.server.utils.BuilderUtils; import org.apache.log4j.Level; import org.apache.log4j.LogManager; import org.apache.log4j.Logger; import org.junit.Assert; import org.junit.Test; public class TestApplicationMasterLauncher { private static final Log LOG = LogFactory .getLog(TestApplicationMasterLauncher.class); private static final class MyContainerManagerImpl implements ContainerManagementProtocol { boolean launched = false; boolean cleanedup = false; String attemptIdAtContainerManager = null; String containerIdAtContainerManager = null; String nmHostAtContainerManager = null; long submitTimeAtContainerManager; int maxAppAttempts; @Override public StartContainersResponse startContainers(StartContainersRequest requests) throws YarnException { StartContainerRequest request = requests.getStartContainerRequests().get(0); LOG.info("Container started by MyContainerManager: " + request); launched = true; Map<String, String> env = request.getContainerLaunchContext().getEnvironment(); Token containerToken = request.getContainerToken(); ContainerTokenIdentifier tokenId = null; try { tokenId = BuilderUtils.newContainerTokenIdentifier(containerToken); } catch (IOException e) { throw RPCUtil.getRemoteException(e); } ContainerId containerId = tokenId.getContainerID(); containerIdAtContainerManager = containerId.toString(); attemptIdAtContainerManager = containerId.getApplicationAttemptId().toString(); nmHostAtContainerManager = tokenId.getNmHostAddress(); submitTimeAtContainerManager = Long.parseLong(env.get(ApplicationConstants.APP_SUBMIT_TIME_ENV)); maxAppAttempts = Integer.parseInt(env.get(ApplicationConstants.MAX_APP_ATTEMPTS_ENV)); return StartContainersResponse.newInstance( new HashMap<String, ByteBuffer>(), new ArrayList<ContainerId>(), new HashMap<ContainerId, SerializedException>()); } @Override public StopContainersResponse stopContainers(StopContainersRequest request) throws YarnException { LOG.info("Container cleaned up by MyContainerManager"); cleanedup = true; return null; } @Override public GetContainerStatusesResponse getContainerStatuses( GetContainerStatusesRequest request) throws YarnException { return null; } } @Test public void testAMLaunchAndCleanup() throws Exception { Logger rootLogger = LogManager.getRootLogger(); rootLogger.setLevel(Level.DEBUG); MyContainerManagerImpl containerManager = new MyContainerManagerImpl(); MockRMWithCustomAMLauncher rm = new MockRMWithCustomAMLauncher( containerManager); rm.start(); MockNM nm1 = rm.registerNode("127.0.0.1:1234", 5120); RMApp app = rm.submitApp(2000); // kick the scheduling nm1.nodeHeartbeat(true); int waitCount = 0; while (containerManager.launched == false && waitCount++ < 20) { LOG.info("Waiting for AM Launch to happen.."); Thread.sleep(1000); } Assert.assertTrue(containerManager.launched); RMAppAttempt attempt = app.getCurrentAppAttempt(); ApplicationAttemptId appAttemptId = attempt.getAppAttemptId(); Assert.assertEquals(appAttemptId.toString(), containerManager.attemptIdAtContainerManager); Assert.assertEquals(app.getSubmitTime(), containerManager.submitTimeAtContainerManager); Assert.assertEquals(app.getRMAppAttempt(appAttemptId) .getMasterContainer().getId() .toString(), containerManager.containerIdAtContainerManager); Assert.assertEquals(nm1.getNodeId().toString(), containerManager.nmHostAtContainerManager); Assert.assertEquals(YarnConfiguration.DEFAULT_RM_AM_MAX_ATTEMPTS, containerManager.maxAppAttempts); MockAM am = new MockAM(rm.getRMContext(), rm .getApplicationMasterService(), appAttemptId); am.registerAppAttempt(); am.unregisterAppAttempt(); //complete the AM container to finish the app normally nm1.nodeHeartbeat(attempt.getAppAttemptId(), 1, ContainerState.COMPLETE); am.waitForState(RMAppAttemptState.FINISHED); waitCount = 0; while (containerManager.cleanedup == false && waitCount++ < 20) { LOG.info("Waiting for AM Cleanup to happen.."); Thread.sleep(1000); } Assert.assertTrue(containerManager.cleanedup); am.waitForState(RMAppAttemptState.FINISHED); rm.stop(); } @SuppressWarnings("unused") @Test(timeout = 100000) public void testallocateBeforeAMRegistration() throws Exception { Logger rootLogger = LogManager.getRootLogger(); boolean thrown = false; rootLogger.setLevel(Level.DEBUG); MockRM rm = new MockRM(); rm.start(); MockNM nm1 = rm.registerNode("h1:1234", 5000); RMApp app = rm.submitApp(2000); // kick the scheduling nm1.nodeHeartbeat(true); RMAppAttempt attempt = app.getCurrentAppAttempt(); MockAM am = rm.sendAMLaunched(attempt.getAppAttemptId()); // request for containers int request = 2; try { AllocateResponse ar = am.allocate("h1", 1000, request, new ArrayList<ContainerId>()); } catch (Exception e) { Assert.assertEquals("Application Master is trying to allocate before " + "registering for: " + attempt.getAppAttemptId().getApplicationId(), e.getMessage()); thrown = true; } // kick the scheduler nm1.nodeHeartbeat(true); try { AllocateResponse amrs = am.allocate(new ArrayList<ResourceRequest>(), new ArrayList<ContainerId>()); } catch (Exception e) { Assert.assertEquals("Application Master is trying to allocate before " + "registering for: " + attempt.getAppAttemptId().getApplicationId(), e.getMessage()); thrown = true; } Assert.assertTrue(thrown); am.registerAppAttempt(); thrown = false; try { am.registerAppAttempt(false); } catch (Exception e) { Assert.assertEquals("Application Master is already registered : " + attempt.getAppAttemptId().getApplicationId(), e.getMessage()); thrown = true; } Assert.assertTrue(thrown); } }