LocalityRouterPolicy extends WeightedRandomRouterPolicy { @Override public SubClusterId getHomeSubcluster( ApplicationSubmissionContext appSubmissionContext, List<SubClusterId> blackListSubClusters) throws YarnException { validate(appSubmissionContext); List<ResourceRequest> rrList = appSubmissionContext.getAMContainerResourceRequests(); if (rrList == null || rrList.isEmpty() || (rrList.size() == 1 && ResourceRequest.isAnyLocation(rrList.get(0).getResourceName()))) { return super .getHomeSubcluster(appSubmissionContext, blackListSubClusters); } if (rrList.size() != 3) { throw new FederationPolicyException( "Invalid number of resource requests: " + rrList.size()); } Map<SubClusterId, SubClusterInfo> activeSubClusters = getActiveSubclusters(); List<SubClusterId> validSubClusters = new ArrayList<>(activeSubClusters.keySet()); FederationPolicyUtils .validateSubClusterAvailability(validSubClusters, blackListSubClusters); if (blackListSubClusters != null) { validSubClusters.removeAll(blackListSubClusters); } try { SubClusterId targetId = null; ResourceRequest nodeRequest = null; ResourceRequest rackRequest = null; ResourceRequest anyRequest = null; for (ResourceRequest rr : rrList) { try { targetId = resolver.getSubClusterForNode(rr.getResourceName()); nodeRequest = rr; } catch (YarnException e) { LOG.error("Cannot resolve node : {}", e.getLocalizedMessage()); } try { resolver.getSubClustersForRack(rr.getResourceName()); rackRequest = rr; } catch (YarnException e) { LOG.error("Cannot resolve rack : {}", e.getLocalizedMessage()); } if (ResourceRequest.isAnyLocation(rr.getResourceName())) { anyRequest = rr; continue; } } if (nodeRequest == null) { throw new YarnException("Missing node request"); } if (rackRequest == null) { throw new YarnException("Missing rack request"); } if (anyRequest == null) { throw new YarnException("Missing any request"); } LOG.info( "Node request: " + nodeRequest.getResourceName() + ", Rack request: " + rackRequest.getResourceName() + ", Any request: " + anyRequest .getResourceName()); if (validSubClusters.contains(targetId) && enabledSCs .contains(targetId)) { LOG.info("Node {} is in SubCluster: {}", nodeRequest.getResourceName(), targetId); return targetId; } else { throw new YarnException("The node " + nodeRequest.getResourceName() + " is in a blacklist SubCluster or not active. "); } } catch (YarnException e) { LOG.error("Validating resource requests failed, Falling back to " + "WeightedRandomRouterPolicy placement: " + e.getMessage()); ResourceRequest amReq = Records.newRecord(ResourceRequest.class); amReq.setPriority(appSubmissionContext.getPriority()); amReq.setResourceName(ResourceRequest.ANY); amReq.setCapability(appSubmissionContext.getResource()); amReq.setNumContainers(1); amReq.setRelaxLocality(true); amReq.setNodeLabelExpression( appSubmissionContext.getNodeLabelExpression()); amReq.setExecutionTypeRequest( ExecutionTypeRequest.newInstance(ExecutionType.GUARANTEED)); appSubmissionContext .setAMContainerResourceRequests(Collections.singletonList(amReq)); return super .getHomeSubcluster(appSubmissionContext, blackListSubClusters); } }  @Override void reinitialize(FederationPolicyInitializationContext policyContext); @Override SubClusterId getHomeSubcluster(
      ApplicationSubmissionContext appSubmissionContext,
      List<SubClusterId> blackListSubClusters); static final Logger LOG; }
@Test public void testNodeInActiveSubCluster() throws YarnException { List<ResourceRequest> requests = new ArrayList<ResourceRequest>(); requests.add(ResourceRequest .newInstance(Priority.UNDEFINED, "node1", Resource.newInstance(10, 1), 1)); requests.add(ResourceRequest .newInstance(Priority.UNDEFINED, "rack1", Resource.newInstance(10, 1), 1)); requests.add(ResourceRequest .newInstance(Priority.UNDEFINED, ResourceRequest.ANY, Resource.newInstance(10, 1), 1)); ApplicationSubmissionContext asc = ApplicationSubmissionContext .newInstance(null, null, null, null, null, false, false, 0, Resources.none(), null, false, null, null); asc.setAMContainerResourceRequests(requests); SubClusterId chosen = ((FederationRouterPolicy) getPolicy()).getHomeSubcluster(asc, null); if (getActiveSubclusters().containsKey( getFederationPolicyContext().getFederationSubclusterResolver() .getSubClusterForNode("node1").getId())) { Assert.assertEquals( getFederationPolicyContext().getFederationSubclusterResolver() .getSubClusterForNode("node1"), chosen); } Assert.assertTrue(getActiveSubclusters().containsKey(chosen)); }
@Test public void testMultipleResourceRequests() throws YarnException { List<ResourceRequest> requests = new ArrayList<ResourceRequest>(); requests.add(ResourceRequest .newInstance(Priority.UNDEFINED, "node1", Resource.newInstance(10, 1), 1)); requests.add(ResourceRequest .newInstance(Priority.UNDEFINED, "node2", Resource.newInstance(10, 1), 1)); ApplicationSubmissionContext asc = ApplicationSubmissionContext .newInstance(null, null, null, null, null, false, false, 0, Resources.none(), null, false, null, null); asc.setAMContainerResourceRequests(requests); try { ((FederationRouterPolicy) getPolicy()).getHomeSubcluster(asc, null); Assert.fail(); } catch (FederationPolicyException e) { Assert.assertTrue( e.getMessage().startsWith("Invalid number of resource requests: ")); } }
@Test public void testNodeNotExists() throws YarnException { List<ResourceRequest> requests = new ArrayList<ResourceRequest>(); boolean relaxLocality = true; requests.add(ResourceRequest .newInstance(Priority.UNDEFINED, "node5", Resource.newInstance(10, 1), 1, relaxLocality)); requests.add(ResourceRequest .newInstance(Priority.UNDEFINED, "rack1", Resource.newInstance(10, 1), 1)); requests.add(ResourceRequest .newInstance(Priority.UNDEFINED, ResourceRequest.ANY, Resource.newInstance(10, 1), 1)); ApplicationSubmissionContext asc = ApplicationSubmissionContext .newInstance(null, null, null, null, null, false, false, 0, Resources.none(), null, false, null, null); asc.setAMContainerResourceRequests(requests); try { ((FederationRouterPolicy) getPolicy()).getHomeSubcluster(asc, null); } catch (FederationPolicyException e) { Assert.fail(); } }
@Test public void testNodeInABlacklistSubCluster() throws YarnException { String subClusterToBlacklist = "subcluster3"; Float value = getPolicyInfo().getRouterPolicyWeights().get(subClusterToBlacklist); getPolicyInfo().getRouterPolicyWeights() .put(new SubClusterIdInfo(subClusterToBlacklist), 0.0f); initializePolicy(new YarnConfiguration()); FederationPoliciesTestUtil .initializePolicyContext(getFederationPolicyContext(), getPolicy(), getPolicyInfo(), getActiveSubclusters(), new Configuration()); List<ResourceRequest> requests = new ArrayList<ResourceRequest>(); boolean relaxLocality = true; requests.add(ResourceRequest .newInstance(Priority.UNDEFINED, "node4", Resource.newInstance(10, 1), 1, relaxLocality)); requests.add(ResourceRequest .newInstance(Priority.UNDEFINED, "rack1", Resource.newInstance(10, 1), 1)); requests.add(ResourceRequest .newInstance(Priority.UNDEFINED, ResourceRequest.ANY, Resource.newInstance(10, 1), 1)); ApplicationSubmissionContext asc = ApplicationSubmissionContext .newInstance(null, null, null, null, null, false, false, 0, Resources.none(), null, false, null, null); asc.setAMContainerResourceRequests(requests); try { SubClusterId targetId = ((FederationRouterPolicy) getPolicy()).getHomeSubcluster(asc, null); Assert.assertNotEquals(targetId.getId(), subClusterToBlacklist); } catch (FederationPolicyException e) { Assert.fail(); } getPolicyInfo().getRouterPolicyWeights() .put(new SubClusterIdInfo(subClusterToBlacklist), value); }
@Test public void testNodeNotInPolicy() throws YarnException { String subClusterToBlacklist = "subcluster3"; Float value = getPolicyInfo().getRouterPolicyWeights().get(subClusterToBlacklist); getPolicyInfo().getRouterPolicyWeights().remove(subClusterToBlacklist); initializePolicy(new YarnConfiguration()); FederationPoliciesTestUtil .initializePolicyContext(getFederationPolicyContext(), getPolicy(), getPolicyInfo(), getActiveSubclusters(), new Configuration()); List<ResourceRequest> requests = new ArrayList<ResourceRequest>(); boolean relaxLocality = true; requests.add(ResourceRequest .newInstance(Priority.UNDEFINED, "node4", Resource.newInstance(10, 1), 1, relaxLocality)); requests.add(ResourceRequest .newInstance(Priority.UNDEFINED, "rack1", Resource.newInstance(10, 1), 1)); requests.add(ResourceRequest .newInstance(Priority.UNDEFINED, ResourceRequest.ANY, Resource.newInstance(10, 1), 1)); ApplicationSubmissionContext asc = ApplicationSubmissionContext .newInstance(null, null, null, null, null, false, false, 0, Resources.none(), null, false, null, null); asc.setAMContainerResourceRequests(requests); try { SubClusterId targetId = ((FederationRouterPolicy) getPolicy()).getHomeSubcluster(asc, null); Assert.assertNotEquals(targetId.getId(), subClusterToBlacklist); } catch (FederationPolicyException e) { Assert.fail(); } getPolicyInfo().getRouterPolicyWeights() .put(new SubClusterIdInfo(subClusterToBlacklist), value); }