LocalityMulticastAMRMProxyPolicy extends AbstractAMRMProxyPolicy { @Override public Map<SubClusterId, List<ResourceRequest>> splitResourceRequests( List<ResourceRequest> resourceRequests, Set<SubClusterId> timedOutSubClusters) throws YarnException { AllocationBookkeeper bookkeeper = new AllocationBookkeeper(); bookkeeper.reinitialize(getActiveSubclusters(), timedOutSubClusters); List<ResourceRequest> nonLocalizedRequests = new ArrayList<ResourceRequest>(); SubClusterId targetId = null; Set<SubClusterId> targetIds = null; for (ResourceRequest rr : resourceRequests) { targetId = null; targetIds = null; if (ResourceRequest.isAnyLocation(rr.getResourceName())) { nonLocalizedRequests.add(rr); continue; } try { targetId = resolver.getSubClusterForNode(rr.getResourceName()); } catch (YarnException e) { } if (bookkeeper.isActiveAndEnabled(targetId)) { bookkeeper.addLocalizedNodeRR(targetId, rr); continue; } try { targetIds = resolver.getSubClustersForRack(rr.getResourceName()); } catch (YarnException e) { } if (targetIds != null && targetIds.size() > 0) { boolean hasActive = false; for (SubClusterId tid : targetIds) { if (bookkeeper.isActiveAndEnabled(tid)) { bookkeeper.addRackRR(tid, rr); hasActive = true; } } if (hasActive) { continue; } } targetId = getSubClusterForUnResolvedRequest(bookkeeper, rr.getAllocationRequestId()); LOG.debug("ERROR resolving sub-cluster for resourceName: {}, picked a " + "random subcluster to forward:{}", rr.getResourceName(), targetId); if (targetIds != null && targetIds.size() > 0) { bookkeeper.addRackRR(targetId, rr); } else { bookkeeper.addLocalizedNodeRR(targetId, rr); } } splitAnyRequests(nonLocalizedRequests, bookkeeper); return bookkeeper.getAnswer(); }  @Override void reinitialize(
      FederationPolicyInitializationContext policyContext); @Override void notifyOfResponse(SubClusterId subClusterId,
      AllocateResponse response); @Override Map<SubClusterId, List<ResourceRequest>> splitResourceRequests(
      List<ResourceRequest> resourceRequests,
      Set<SubClusterId> timedOutSubClusters); static final Logger LOG; }
@Test(timeout = 5000) public void testStressPolicy() throws Exception { getPolicyInfo().setHeadroomAlpha(1.0f); initializePolicy(); addHomeSubClusterAsActive(); int numRR = 1000; List<ResourceRequest> resourceRequests = createLargeRandomList(numRR); prepPolicyWithHeadroom(true); int numIterations = 1000; long tstart = System.currentTimeMillis(); for (int i = 0; i < numIterations; i++) { Map<SubClusterId, List<ResourceRequest>> response = ((FederationAMRMProxyPolicy) getPolicy()).splitResourceRequests( resourceRequests, new HashSet<SubClusterId>()); validateSplit(response, resourceRequests); } long tend = System.currentTimeMillis(); LOG.info("Performed " + numIterations + " policy invocations (and " + "validations) in " + (tend - tstart) + "ms"); }
@Test public void testFWDAllZeroANY() throws Exception { getPolicyInfo().setHeadroomAlpha(0.5f); initializePolicy(); List<ResourceRequest> resourceRequests = createZeroSizedANYRequest(); prepPolicyWithHeadroom(true); Map<SubClusterId, List<ResourceRequest>> response = ((FederationAMRMProxyPolicy) getPolicy()).splitResourceRequests( resourceRequests, new HashSet<SubClusterId>()); prettyPrintRequests(response); validateSplit(response, resourceRequests); checkExpectedAllocation(response, "subcluster0", 1, 0); checkExpectedAllocation(response, "subcluster1", 1, 0); checkExpectedAllocation(response, "subcluster2", 1, 0); checkExpectedAllocation(response, "subcluster3", -1, -1); checkExpectedAllocation(response, "subcluster4", -1, -1); checkExpectedAllocation(response, "subcluster5", -1, -1); checkTotalContainerAllocation(response, 0); }
@Test public void testSplitBasedOnHeadroomAndWeights() throws Exception { getPolicyInfo().setHeadroomAlpha(0.5f); initializePolicy(); List<ResourceRequest> resourceRequests = createSimpleRequest(); prepPolicyWithHeadroom(true); Map<SubClusterId, List<ResourceRequest>> response = ((FederationAMRMProxyPolicy) getPolicy()).splitResourceRequests( resourceRequests, new HashSet<SubClusterId>()); prettyPrintRequests(response); validateSplit(response, resourceRequests); checkExpectedAllocation(response, "subcluster0", 1, 42); checkExpectedAllocation(response, "subcluster1", 1, 12); checkExpectedAllocation(response, "subcluster2", 1, 20); checkExpectedAllocation(response, "subcluster3", -1, -1); checkExpectedAllocation(response, "subcluster4", -1, -1); checkExpectedAllocation(response, "subcluster5", 1, 25); checkTotalContainerAllocation(response, 100); }
@Test public void testSplitAllocateRequest() throws Exception { initializePolicy(); addHomeSubClusterAsActive(); FederationPoliciesTestUtil.initializePolicyContext( getFederationPolicyContext(), getPolicy(), getPolicyInfo(), getActiveSubclusters(), new Configuration()); List<ResourceRequest> resourceRequests = createComplexRequest(); Map<SubClusterId, List<ResourceRequest>> response = ((FederationAMRMProxyPolicy) getPolicy()).splitResourceRequests( resourceRequests, new HashSet<SubClusterId>()); validateSplit(response, resourceRequests); prettyPrintRequests(response); checkExpectedAllocation(response, getHomeSubCluster().getId(), 7, 29); checkExpectedAllocation(response, "subcluster0", 10, 32); checkExpectedAllocation(response, "subcluster1", 5, 26); checkExpectedAllocation(response, "subcluster2", 4, 23); checkExpectedAllocation(response, "subcluster3", -1, -1); checkExpectedAllocation(response, "subcluster4", -1, -1); checkExpectedAllocation(response, "subcluster5", 1, 20); checkTotalContainerAllocation(response, 130); for (ResourceRequest rr : response.get(getHomeSubCluster())) { Assert.assertTrue( rr.getAllocationRequestId() == 2L || rr.getAllocationRequestId() == 4L || rr.getAllocationRequestId() == 5L); } List<ResourceRequest> rrs = response.get(SubClusterId.newInstance("subcluster0")); for (ResourceRequest rr : rrs) { Assert.assertTrue(rr.getAllocationRequestId() != 1L); Assert.assertTrue(rr.getAllocationRequestId() != 4L); } for (ResourceRequest rr : response .get(SubClusterId.newInstance("subcluster1"))) { Assert.assertTrue(rr.getAllocationRequestId() == 1L || rr.getAllocationRequestId() == 2L); } for (ResourceRequest rr : response .get(SubClusterId.newInstance("subcluster2"))) { Assert.assertTrue(rr.getAllocationRequestId() == 1L || rr.getAllocationRequestId() == 2L); } for (ResourceRequest rr : response .get(SubClusterId.newInstance("subcluster5"))) { Assert.assertTrue(rr.getAllocationRequestId() == 2); Assert.assertTrue(rr.getRelaxLocality()); } }
@Test public void testCancelWithLocalizedResource() throws YarnException { getPolicyInfo().setHeadroomAlpha(1.0f); initializePolicy(); List<ResourceRequest> resourceRequests = new ArrayList<>(); prepPolicyWithHeadroom(true); resourceRequests.add(FederationPoliciesTestUtil.createResourceRequest(0L, "subcluster0-rack0-host0", 1024, 1, 1, 1, null, false)); resourceRequests.add(FederationPoliciesTestUtil.createResourceRequest(0L, "subcluster0-rack0", 1024, 1, 1, 1, null, false)); resourceRequests.add(FederationPoliciesTestUtil.createResourceRequest(0L, ResourceRequest.ANY, 1024, 1, 1, 0, null, false)); Map<SubClusterId, List<ResourceRequest>> response = ((FederationAMRMProxyPolicy) getPolicy()).splitResourceRequests( resourceRequests, new HashSet<SubClusterId>()); checkExpectedAllocation(response, "subcluster0", 3, 1); checkExpectedAllocation(response, "subcluster1", 1, 0); checkExpectedAllocation(response, "subcluster2", 1, 0); checkExpectedAllocation(response, "subcluster3", -1, -1); checkExpectedAllocation(response, "subcluster4", -1, -1); checkExpectedAllocation(response, "subcluster5", -1, -1); resourceRequests.clear(); resourceRequests.add(FederationPoliciesTestUtil.createResourceRequest(0L, "subcluster0-rack0-host0", 1024, 1, 1, 0, null, false)); resourceRequests.add(FederationPoliciesTestUtil.createResourceRequest(0L, "subcluster0-rack0", 1024, 1, 1, 0, null, false)); resourceRequests.add(FederationPoliciesTestUtil.createResourceRequest(0L, ResourceRequest.ANY, 1024, 1, 1, 100, null, false)); response = ((FederationAMRMProxyPolicy) getPolicy()) .splitResourceRequests(resourceRequests, new HashSet<SubClusterId>()); checkExpectedAllocation(response, "subcluster0", 3, 60); checkExpectedAllocation(response, "subcluster1", 1, -1); checkExpectedAllocation(response, "subcluster2", 1, 15); checkExpectedAllocation(response, "subcluster5", 1, 25); checkTotalContainerAllocation(response, 100); }
@Test public void testSubClusterExpiry() throws Exception { getPolicyInfo().setHeadroomAlpha(1.0f); YarnConfiguration conf = new YarnConfiguration(); conf.setLong(YarnConfiguration.FEDERATION_AMRMPROXY_SUBCLUSTER_TIMEOUT, 500); initializePolicy(conf); List<ResourceRequest> resourceRequests = createSimpleRequest(); prepPolicyWithHeadroom(true); Set<SubClusterId> expiredSCList = new HashSet<>(); Map<SubClusterId, List<ResourceRequest>> response = ((FederationAMRMProxyPolicy) getPolicy()) .splitResourceRequests(resourceRequests, expiredSCList); prettyPrintRequests(response); validateSplit(response, resourceRequests); checkExpectedAllocation(response, "subcluster0", 1, 60); checkExpectedAllocation(response, "subcluster1", 1, -1); checkExpectedAllocation(response, "subcluster2", 1, 15); checkExpectedAllocation(response, "subcluster5", 1, 25); checkTotalContainerAllocation(response, 100); Thread.sleep(800); expiredSCList.add(SubClusterId.newInstance("subcluster0")); expiredSCList.add(SubClusterId.newInstance("subcluster5")); response = ((FederationAMRMProxyPolicy) getPolicy()) .splitResourceRequests(resourceRequests, expiredSCList); prettyPrintRequests(response); validateSplit(response, resourceRequests); checkExpectedAllocation(response, "subcluster0", 1, -1); checkExpectedAllocation(response, "subcluster1", 1, -1); checkExpectedAllocation(response, "subcluster2", 1, 100); checkExpectedAllocation(response, "subcluster5", 1, -1); checkTotalContainerAllocation(response, 100); }