LocalityMulticastAMRMProxyPolicy extends AbstractAMRMProxyPolicy { @VisibleForTesting protected ArrayList<Integer> computeIntegerAssignment(int totalNum, ArrayList<Float> weightsList) throws YarnException { int i, residue; ArrayList<Integer> ret = new ArrayList<>(); float totalWeight = 0, totalNumFloat = totalNum; if (weightsList.size() == 0) { return ret; } for (i = 0; i < weightsList.size(); i++) { ret.add(0); if (weightsList.get(i) > 0) { totalWeight += weightsList.get(i); } } if (totalWeight == 0) { StringBuilder sb = new StringBuilder(); for (Float weight : weightsList) { sb.append(weight + ", "); } throw new FederationPolicyException( "No positive value found in weight array " + sb.toString()); } residue = totalNum; for (i = 0; i < weightsList.size(); i++) { if (weightsList.get(i) > 0) { int base = (int) (totalNumFloat * weightsList.get(i) / totalWeight); ret.set(i, ret.get(i) + base); residue -= base; } } for (i = 0; i < residue; i++) { int index = FederationPolicyUtils.getWeightedRandom(weightsList); ret.set(index, ret.get(index) + 1); } return ret; }  @Override void reinitialize(
      FederationPolicyInitializationContext policyContext); @Override void notifyOfResponse(SubClusterId subClusterId,
      AllocateResponse response); @Override Map<SubClusterId, List<ResourceRequest>> splitResourceRequests(
      List<ResourceRequest> resourceRequests); static final Logger LOG; }
@Test public void testIntegerAssignment() throws YarnException { float[] weights = new float[] {0, 0.1f, 0.2f, 0.2f, -0.1f, 0.1f, 0.2f, 0.1f, 0.1f}; int[] expectedMin = new int[] {0, 1, 3, 3, 0, 1, 3, 1, 1}; ArrayList<Float> weightsList = new ArrayList<>(); for (float weight : weights) { weightsList.add(weight); } LocalityMulticastAMRMProxyPolicy policy = (LocalityMulticastAMRMProxyPolicy) getPolicy(); for (int i = 0; i < 500000; i++) { ArrayList<Integer> allocations = policy.computeIntegerAssignment(19, weightsList); int sum = 0; for (int j = 0; j < weights.length; j++) { sum += allocations.get(j); if (allocations.get(j) < expectedMin[j]) { Assert.fail(allocations.get(j) + " at index " + j + " should be at least " + expectedMin[j] + ". Allocation array: " + printList(allocations)); } } Assert.assertEquals( "Expect sum to be 19 in array: " + printList(allocations), 19, sum); } }