ApplicationMasterService extends AbstractService implements
    ApplicationMasterProtocol { @Override public AllocateResponse allocate(AllocateRequest request) throws YarnException, IOException { AMRMTokenIdentifier amrmTokenIdentifier = YarnServerSecurityUtils.authorizeRequest(); ApplicationAttemptId appAttemptId = amrmTokenIdentifier.getApplicationAttemptId(); this.amLivelinessMonitor.receivedPing(appAttemptId); AllocateResponseLock lock = responseMap.get(appAttemptId); if (lock == null) { String message = "Application attempt " + appAttemptId + " doesn't exist in ApplicationMasterService cache."; LOG.error(message); throw new ApplicationAttemptNotFoundException(message); } synchronized (lock) { AllocateResponse lastResponse = lock.getAllocateResponse(); if (!hasApplicationMasterRegistered(appAttemptId)) { String message = "AM is not registered for known application attempt: " + appAttemptId + " or RM had restarted after AM registered. " + " AM should re-register."; throw new ApplicationMasterNotRegisteredException(message); } if (AMRMClientUtils.getNextResponseId( request.getResponseId()) == lastResponse.getResponseId()) { return lastResponse; } else if (request.getResponseId() != lastResponse.getResponseId()) { throw new InvalidApplicationMasterRequestException(AMRMClientUtils .assembleInvalidResponseIdExceptionMessage(appAttemptId, lastResponse.getResponseId(), request.getResponseId())); } AllocateResponse response = recordFactory.newRecordInstance(AllocateResponse.class); this.amsProcessingChain.allocate( amrmTokenIdentifier.getApplicationAttemptId(), request, response); MasterKeyData nextMasterKey = this.rmContext.getAMRMTokenSecretManager().getNextMasterKeyData(); if (nextMasterKey != null && nextMasterKey.getMasterKey().getKeyId() != amrmTokenIdentifier .getKeyId()) { RMApp app = this.rmContext.getRMApps().get(appAttemptId.getApplicationId()); RMAppAttempt appAttempt = app.getRMAppAttempt(appAttemptId); RMAppAttemptImpl appAttemptImpl = (RMAppAttemptImpl)appAttempt; Token<AMRMTokenIdentifier> amrmToken = appAttempt.getAMRMToken(); if (nextMasterKey.getMasterKey().getKeyId() != appAttemptImpl.getAMRMTokenKeyId()) { LOG.info("The AMRMToken has been rolled-over. Send new AMRMToken back" + " to application: " + appAttemptId.getApplicationId()); amrmToken = rmContext.getAMRMTokenSecretManager() .createAndGetAMRMToken(appAttemptId); appAttemptImpl.setAMRMToken(amrmToken); } response.setAMRMToken(org.apache.hadoop.yarn.api.records.Token .newInstance(amrmToken.getIdentifier(), amrmToken.getKind() .toString(), amrmToken.getPassword(), amrmToken.getService() .toString())); } response.setResponseId( AMRMClientUtils.getNextResponseId(lastResponse.getResponseId())); lock.setAllocateResponse(response); return response; } } ApplicationMasterService(RMContext rmContext,
      YarnScheduler scheduler);  ApplicationMasterService(String name, RMContext rmContext,
      YarnScheduler scheduler); @Private InetSocketAddress getBindAddress(); @Override RegisterApplicationMasterResponse registerApplicationMaster(
      RegisterApplicationMasterRequest request); @Override FinishApplicationMasterResponse finishApplicationMaster(
      FinishApplicationMasterRequest request); boolean hasApplicationMasterRegistered(
      ApplicationAttemptId appAttemptId); @Override AllocateResponse allocate(AllocateRequest request); void registerAppAttempt(ApplicationAttemptId attemptId); void unregisterAttempt(ApplicationAttemptId attemptId); void refreshServiceAcls(Configuration configuration, 
      PolicyProvider policyProvider); @VisibleForTesting Server getServer();  }
@Test public void testValidateRequestCapacityAgainstMinMaxAllocationWithDifferentUnits() throws Exception { Map<String, ResourceInformation> riMap = new HashMap<>(); ResourceInformation memory = ResourceInformation.newInstance(ResourceInformation.MEMORY_MB.getName(), ResourceInformation.MEMORY_MB.getUnits(), YarnConfiguration.DEFAULT_RM_SCHEDULER_MINIMUM_ALLOCATION_MB, YarnConfiguration.DEFAULT_RM_SCHEDULER_MAXIMUM_ALLOCATION_MB); ResourceInformation vcores = ResourceInformation.newInstance(ResourceInformation.VCORES.getName(), ResourceInformation.VCORES.getUnits(), YarnConfiguration.DEFAULT_RM_SCHEDULER_MINIMUM_ALLOCATION_VCORES, DEFAULT_RM_SCHEDULER_MAXIMUM_ALLOCATION_VCORES); ResourceInformation res1 = ResourceInformation.newInstance("res_1", "G", 0, 4); riMap.put(ResourceInformation.MEMORY_URI, memory); riMap.put(ResourceInformation.VCORES_URI, vcores); riMap.put("res_1", res1); ResourceUtils.initializeResourcesFromResourceInformationMap(riMap); FairSchedulerConfiguration fsConf = new FairSchedulerConfiguration(); YarnConfiguration yarnConf = new YarnConfiguration(fsConf); yarnConf.setBoolean(TestResourceProfiles.TEST_CONF_RESET_RESOURCE_TYPES, false); yarnConf.setClass(YarnConfiguration.RM_SCHEDULER, FairScheduler.class, ResourceScheduler.class); yarnConf.setBoolean(YarnConfiguration.RM_RESOURCE_PROFILES_ENABLED, false); MockRM rm = new MockRM(yarnConf); rm.start(); MockNM nm1 = rm.registerNode("199.99.99.1:1234", ResourceTypesTestHelper.newResource( DEFAULT_RM_SCHEDULER_MAXIMUM_ALLOCATION_MB, DEFAULT_RM_SCHEDULER_MAXIMUM_ALLOCATION_VCORES, ImmutableMap.<String, String> builder() .put("res_1", "5G").build())); RMApp app1 = rm.submitApp(GB, "app", "user", null, "default"); MockAM am1 = MockRM.launchAndRegisterAM(app1, rm, nm1); try { am1.allocate(Collections.singletonList(ResourceRequest.newBuilder() .capability(ResourceTypesTestHelper.newResource(4 * GB, 1, ImmutableMap.<String, String> builder() .put("res_1", "500M") .build())) .numContainers(1).resourceName("*").build()), null); } catch (InvalidResourceRequestException e) { fail( "Allocate request should be accepted but exception was thrown: " + e); } rm.close(); }
@Test(timeout = 300000) public void testValidateRequestCapacityAgainstMinMaxAllocationFor3rdResourceTypes() throws Exception { Map<String, ResourceInformation> riMap = new HashMap<>(); ResourceInformation memory = ResourceInformation.newInstance( ResourceInformation.MEMORY_MB.getName(), ResourceInformation.MEMORY_MB.getUnits(), YarnConfiguration.DEFAULT_RM_SCHEDULER_MINIMUM_ALLOCATION_MB, DEFAULT_RM_SCHEDULER_MAXIMUM_ALLOCATION_MB); ResourceInformation vcores = ResourceInformation.newInstance( ResourceInformation.VCORES.getName(), ResourceInformation.VCORES.getUnits(), YarnConfiguration.DEFAULT_RM_SCHEDULER_MINIMUM_ALLOCATION_VCORES, DEFAULT_RM_SCHEDULER_MAXIMUM_ALLOCATION_VCORES); ResourceInformation res1 = ResourceInformation.newInstance("res_1", ResourceInformation.VCORES.getUnits(), 0, 4); riMap.put(ResourceInformation.MEMORY_URI, memory); riMap.put(ResourceInformation.VCORES_URI, vcores); riMap.put("res_1", res1); ResourceUtils.initializeResourcesFromResourceInformationMap(riMap); CapacitySchedulerConfiguration csconf = new CapacitySchedulerConfiguration(); csconf.setResourceComparator(DominantResourceCalculator.class); YarnConfiguration yarnConf = new YarnConfiguration(csconf); yarnConf.setBoolean(TestResourceProfiles.TEST_CONF_RESET_RESOURCE_TYPES, false); yarnConf.setClass(YarnConfiguration.RM_SCHEDULER, CapacityScheduler.class, ResourceScheduler.class); yarnConf.setBoolean(YarnConfiguration.RM_RESOURCE_PROFILES_ENABLED, false); MockRM rm = new MockRM(yarnConf); rm.start(); CapacityScheduler cs = (CapacityScheduler) rm.getResourceScheduler(); LeafQueue leafQueue = (LeafQueue) cs.getQueue("default"); MockNM nm1 = rm.registerNode("199.99.99.1:1234", TestUtils .createResource(DEFAULT_RM_SCHEDULER_MAXIMUM_ALLOCATION_MB, DEFAULT_RM_SCHEDULER_MAXIMUM_ALLOCATION_VCORES, ImmutableMap.of("res_1", 4))); RMApp app1 = rm.submitApp(GB, "app", "user", null, "default"); MockAM am1 = MockRM.launchAndRegisterAM(app1, rm, nm1); Assert.assertEquals(Resource.newInstance(GB, 1), leafQueue.getUsedResources()); boolean exception = false; try { am1.allocate(Collections.singletonList(ResourceRequest.newBuilder() .capability(TestUtils.createResource(9 * GB, 1, ImmutableMap.of("res_1", 1))) .numContainers(1) .resourceName("*") .build()), null); } catch (InvalidResourceRequestException e) { exception = true; } Assert.assertTrue(exception); exception = false; try { am1.allocate(Collections.singletonList(ResourceRequest.newBuilder() .capability( TestUtils.createResource(8 * GB, 18, ImmutableMap.of("res_1", 1))) .numContainers(1) .resourceName("*") .build()), null); } catch (InvalidResourceRequestException e) { exception = true; } Assert.assertTrue(exception); exception = false; try { am1.allocate(Collections.singletonList(ResourceRequest.newBuilder() .capability(TestUtils.createResource(8 * GB, 1, ImmutableMap.of("res_1", 100))) .numContainers(1) .resourceName("*") .build()), null); } catch (InvalidResourceRequestException e) { exception = true; } Assert.assertTrue(exception); rm.close(); }