Support placement constraint

tony-core/src/main/java/com/linkedin/tony/Constants.java

@@ -158,6 +158,9 @@ public class Constants {
   public static final String VCORES = "vcores";
   public static final String GPUS = "gpus";
 
+  public static final String ALLOCATION_TAGS = "allocation-tags";
+  public static final String PLACEMENT_SPEC = "placement-spec";
+
   // pid environment variable set by YARN
   public static final String JVM_PID = "JVM_PID";
 

tony-core/src/main/java/com/linkedin/tony/HadoopCompatibleAdapter.java

@@ -18,21 +18,37 @@
 import java.lang.reflect.Constructor;
 import java.lang.reflect.InvocationTargetException;
 import java.lang.reflect.Method;
+import java.util.ArrayList;
+import java.util.Arrays;
+import java.util.Collections;
 import java.util.HashMap;
+import java.util.HashSet;
+import java.util.List;
 import java.util.Map;
+import java.util.Set;
+import java.util.concurrent.atomic.AtomicLong;
 
-import org.apache.commons.lang.StringUtils;
+import com.linkedin.tony.models.JobContainerRequest;
+import com.linkedin.tony.util.Utils;
+import org.apache.commons.collections.CollectionUtils;
+import org.apache.commons.lang3.StringUtils;
 import org.apache.commons.logging.Log;
 import org.apache.commons.logging.LogFactory;
 import org.apache.hadoop.conf.Configuration;
 import org.apache.hadoop.yarn.api.records.Container;
+import org.apache.hadoop.yarn.api.records.ExecutionTypeRequest;
+import org.apache.hadoop.yarn.api.records.Priority;
 import org.apache.hadoop.yarn.api.records.Resource;
 import org.apache.hadoop.yarn.api.records.ResourceInformation;
+import org.apache.hadoop.yarn.client.api.AMRMClient;
+import org.apache.hadoop.yarn.client.api.async.AMRMClientAsync;
 import org.apache.hadoop.yarn.util.resource.ResourceUtils;
 
 public final class HadoopCompatibleAdapter {
     private static final Log LOG = LogFactory.getLog(HadoopCompatibleAdapter.class);
 
+    private static final AtomicLong ALLOCATE_ID_COUNTER = new AtomicLong(1);
+
     private HadoopCompatibleAdapter() {
 
     }
@@ -157,4 +173,61 @@ public static boolean existGPUResource() {
             return false;
         }
     }
+
+    public static void constructAndAddSchedulingRequest(AMRMClientAsync<AMRMClient.ContainerRequest> amRMClient,
+            JobContainerRequest containerRequest) {
+        try {
+            List<Object> reqs = new ArrayList<>();
+            Object schedReq = constructSchedulingRequest(containerRequest);
+            LOG.info("Request schedling containers ask: " + schedReq);
+            for (int i = 0; i < containerRequest.getNumInstances(); i++) {
+                reqs.add(schedReq);
+            }
+            Method addMethod = Arrays.stream(amRMClient.getClass().getMethods())
+                    .filter(x -> x.getName().equals("addSchedulingRequests") && x.getParameterCount() == 1)
+                    .findFirst().get();
+            addMethod.invoke(amRMClient, reqs);
+        } catch (Exception e) {
+            throw new RuntimeException("Errors on adding scheduing request.", e);
+        }
+    }
+
+    private static Object constructSchedulingRequest(JobContainerRequest containerRequest) {
+        try {
+            Priority priority = Priority.newInstance(containerRequest.getPriority());
+            Resource capability = Resource.newInstance((int) containerRequest.getMemory(), containerRequest.getVCores());
+            if (containerRequest.getGPU() > 0) {
+                Utils.setCapabilityGPU(capability, containerRequest.getGPU());
+            }
+            Set<String> allocationTags = CollectionUtils.isEmpty(containerRequest.getAllocationTags())
+                    ? Collections.singleton("") : new HashSet<>(containerRequest.getAllocationTags());
+
+            Class<?> placementConstraintCls =
+                    Class.forName("org.apache.hadoop.yarn.util.constraint.PlacementConstraintParser");
+            Method parseMethod = placementConstraintCls.getMethod("parseExpression", String.class);
+
+            Object parsedObj = parseMethod.invoke(placementConstraintCls, containerRequest.getPlacementSpec());
+            Class<?> abstractConstraintCls =
+                    Class.forName("org.apache.hadoop.yarn.api.resource.PlacementConstraint$AbstractConstraint");
+
+            Object placementConstraintObj = abstractConstraintCls.getMethod("build").invoke(parsedObj);
+
+            Class<?> resourceSizingCls = Class.forName("org.apache.hadoop.yarn.api.records.ResourceSizing");
+            Method resourceSizingMethod = Arrays.stream(resourceSizingCls.getMethods())
+                    .filter(x -> x.getName().equals("newInstance") && x.getParameterCount() == 1).findFirst().get();
+            Object resourceSizingObj = resourceSizingMethod.invoke(null, capability);
+
+            Class<?> schedulingReqCls = Class.forName("org.apache.hadoop.yarn.api.records.SchedulingRequest");
+            Method newInstanceMethod = Arrays.stream(schedulingReqCls.getMethods())
+                    .filter(x -> x.getName().equals("newInstance") && x.getParameterCount() == 6).findFirst().get();
+
+            Object schedReq = newInstanceMethod.invoke(null, ALLOCATE_ID_COUNTER.incrementAndGet(), priority,
+                    ExecutionTypeRequest.newInstance(), allocationTags,
+                    resourceSizingObj, placementConstraintObj);
+
+            return schedReq;
+        } catch (ClassNotFoundException | NoSuchMethodException | IllegalAccessException | InvocationTargetException e) {
+            throw new RuntimeException("Errors on constructing scheduling requests of Yarn.", e);
+        }
+    }
 }

tony-core/src/main/java/com/linkedin/tony/TaskScheduler.java

@@ -36,7 +36,6 @@ public class TaskScheduler {
   // job with dependency -> (dependent job name, number of instances for that job)
   private Map<JobContainerRequest, Map<String, Integer>> taskDependencyMap = new HashMap<>();
   private Map<String, LocalResource> localResources;
-  private Map<String, List<AMRMClient.ContainerRequest>> jobTypeToContainerRequestsMap = new HashMap<>();
   private Map<String, Map<String, LocalResource>> jobTypeToContainerResources;
 
   boolean dependencyCheckPassed = true;
@@ -90,16 +89,20 @@ boolean checkDependencySatisfied(JobContainerRequest request) {
   }
 
   private void scheduleJob(JobContainerRequest request) {
-    AMRMClient.ContainerRequest containerAsk = Utils.setupContainerRequestForRM(request);
-    String jobName = request.getJobName();
-    if (!jobTypeToContainerRequestsMap.containsKey(jobName)) {
-      jobTypeToContainerRequestsMap.put(jobName, new ArrayList<>());
-      jobTypeToContainerResources.put(jobName, getContainerResources(jobName));
-    }
-    jobTypeToContainerRequestsMap.get(request.getJobName()).add(containerAsk);
-    for (int i = 0; i < request.getNumInstances(); i++) {
-      amRMClient.addContainerRequest(containerAsk);
+    if (request.getPlacementSpec() != null) {
+      // this should use newer api of Yarn with this placement constraint feature,
+      // only be supported in hadoop 3.2.x
+      HadoopCompatibleAdapter.constructAndAddSchedulingRequest(amRMClient, request);
+    } else {
+      AMRMClient.ContainerRequest containerAsk = Utils.setupContainerRequestForRM(request);
+      for (int i = 0; i < request.getNumInstances(); i++) {
+        amRMClient.addContainerRequest(containerAsk);
+      }
     }
+
+    String jobName = request.getJobName();
+    jobTypeToContainerResources.putIfAbsent(jobName, getContainerResources(jobName));
+
     session.addNumExpectedTask(request.getNumInstances());
   }
 

tony-core/src/main/java/com/linkedin/tony/TonyClient.java

@@ -53,7 +53,7 @@
 import org.apache.commons.cli.Options;
 import org.apache.commons.cli.ParseException;
 import org.apache.commons.io.FileUtils;
-import org.apache.commons.lang.StringUtils;
+import org.apache.commons.lang3.StringUtils;
 import org.apache.commons.logging.Log;
 import org.apache.commons.logging.LogFactory;
 import org.apache.hadoop.conf.Configuration;

tony-core/src/main/java/com/linkedin/tony/TonyConfigurationKeys.java

@@ -241,6 +241,22 @@ public static String getResourcesKey(String jobName) {
     return String.format(TONY_PREFIX + "%s.resources", jobName);
   }
 
+  public static String getPlacementSpecKey(String jobName) {
+    return String.format(TONY_PREFIX + "%s.placement-spec", jobName);
+  }
+
+  public static String getAllocationSpecKey(String jobName) {
+    return String.format(TONY_PREFIX + "%s.allocation-tags", jobName);
+  }
+
+  public static String getContainerPlacementSpecKey() {
+    return TONY_PREFIX + "containers.placement-spec";
+  }
+
+  public static String getContainerAllocationTagsKey() {
+    return TONY_PREFIX + "containers.allocation-tags";
+  }
+
   // Resources for all containers
   public static String getContainerResourcesKey() {
     return TONY_PREFIX + "containers.resources";

tony-core/src/main/java/com/linkedin/tony/TonySession.java

@@ -21,7 +21,7 @@
 import java.util.concurrent.ConcurrentHashMap;
 import java.util.stream.Collectors;
 
-import org.apache.commons.lang.StringUtils;
+import org.apache.commons.lang3.StringUtils;
 import org.apache.commons.logging.Log;
 import org.apache.commons.logging.LogFactory;
 import org.apache.hadoop.conf.Configuration;

tony-core/src/main/java/com/linkedin/tony/models/JobContainerRequest.java

@@ -17,6 +17,25 @@ public class JobContainerRequest {
   private String nodeLabelsExpression;
   private List<String> dependsOn;
 
+  private String placementSpec;
+  private List<String> allocationTags;
+
+  public JobContainerRequest(String jobName, int numInstances, long memory, int vCores, int gpu, int priority,
+          String nodeLabelsExpression, final List<String> dependsOn, String placementSpec,
+          List<String> allocationTags) {
+    this.numInstances = numInstances;
+    this.memory = memory;
+    this.vCores = vCores;
+    this.priority = priority;
+    this.gpu = gpu;
+    this.jobName = jobName;
+    this.nodeLabelsExpression = nodeLabelsExpression;
+    this.dependsOn = dependsOn;
+    this.placementSpec = placementSpec;
+    this.allocationTags = allocationTags;
+  }
+
+
   public JobContainerRequest(String jobName, int numInstances, long memory, int vCores, int gpu, int priority,
       String nodeLabelsExpression, final List<String> dependsOn) {
     this.numInstances = numInstances;
@@ -60,4 +79,12 @@ public String getNodeLabelsExpression() {
   public final List<String> getDependsOn() {
     return dependsOn;
   }
+
+  public String getPlacementSpec() {
+    return placementSpec;
+  }
+
+  public List<String> getAllocationTags() {
+    return allocationTags;
+  }
 }

tony-core/src/main/java/com/linkedin/tony/runtime/HorovodRuntime.java

@@ -25,7 +25,7 @@
 import java.util.Map;
 import java.util.stream.Collectors;
 
-import org.apache.commons.lang.StringUtils;
+import org.apache.commons.lang3.StringUtils;
 import org.apache.hadoop.conf.Configuration;
 import org.apache.hadoop.yarn.api.records.FinalApplicationStatus;
 

tony-core/src/main/java/com/linkedin/tony/runtime/MLGenericRuntime.java

@@ -23,7 +23,7 @@
 import java.util.stream.Collectors;
 
 import org.apache.commons.collections.CollectionUtils;
-import org.apache.commons.lang.StringUtils;
+import org.apache.commons.lang3.StringUtils;
 import org.apache.commons.lang3.tuple.Pair;
 import org.apache.hadoop.conf.Configuration;
 import org.apache.hadoop.yarn.api.records.FinalApplicationStatus;

tony-core/src/main/java/com/linkedin/tony/util/Utils.java

@@ -409,6 +409,17 @@ public static Map<String, JobContainerRequest> parseContainerRequests(Configurat
               TonyConfigurationKeys.DEFAULT_VCORES);
       int gpus = conf.getInt(TonyConfigurationKeys.getResourceKey(jobName, Constants.GPUS),
               TonyConfigurationKeys.DEFAULT_GPUS);
+
+      String placementSpec = conf.get(
+              TonyConfigurationKeys.getPlacementSpecKey(jobName),
+              conf.get(TonyConfigurationKeys.getContainerPlacementSpecKey())
+      );
+      String[] allocationTagsArr = conf.getStrings(
+              TonyConfigurationKeys.getAllocationSpecKey(jobName),
+              conf.getStrings(TonyConfigurationKeys.getContainerAllocationTagsKey())
+      );
+      List<String> allocationTags = allocationTagsArr == null ? null : new ArrayList<>(Arrays.asList(allocationTagsArr));
+
       if (gpus > 0 && !HadoopCompatibleAdapter.existGPUResource()) {
         throw new RuntimeException(String.format("User requested %d GPUs for job '%s' but GPU is not available on the cluster. ",
             gpus, jobName));
@@ -431,7 +442,7 @@ public static Map<String, JobContainerRequest> parseContainerRequests(Configurat
         // We rely on unique priority behavior to match allocation request to task in Hadoop 2.7
         containerRequests.put(jobName,
                 new JobContainerRequest(jobName, numInstances, memory, vCores, gpus, priority,
-                    nodeLabel, dependsOn));
+                    nodeLabel, dependsOn, placementSpec, allocationTags));
         priority++;
       }
     }

tony-core/src/test/java/com/linkedin/tony/TestPortAllocation.java

@@ -12,7 +12,7 @@
 import java.io.IOException;
 import java.net.BindException;
 import java.time.Duration;
-import org.apache.commons.lang.SystemUtils;
+import org.apache.commons.lang3.SystemUtils;
 import org.testng.annotations.Test;
 
 public class TestPortAllocation {

tony-core/src/test/java/com/linkedin/tony/util/TestUtils.java

@@ -120,6 +120,39 @@ public void testParseContainerRequestsShouldFail() {
     Utils.parseContainerRequests(conf);
   }
 
+  @Test
+  public void testParsePlacementSpecAndAllocationTags() {
+    Configuration conf = new Configuration();
+    conf.addResource("tony-default.xml");
+    conf.setInt("tony.worker.instances", 1);
+
+    // case1: set nothing
+    Map<String, JobContainerRequest> containerRequests = Utils.parseContainerRequests(conf);
+    assertNull(containerRequests.get("worker").getPlacementSpec());
+    assertNull(containerRequests.get("worker").getAllocationTags());
+
+    // case2: set all
+    conf = new Configuration();
+    conf.addResource("tony-default.xml");
+    conf.setInt("tony.worker.instances", 1);
+    conf.setStrings("tony.worker.placement-spec", "java=true");
+    conf.setStrings("tony.worker.allocation-tags", "tony");
+    containerRequests = Utils.parseContainerRequests(conf);
+
+    assertEquals("java=true", containerRequests.get("worker").getPlacementSpec());
+    assertEquals(Arrays.asList("tony"), containerRequests.get("worker").getAllocationTags());
+
+    // case3: set nothing for job, but it will fallback to container setting
+    conf = new Configuration();
+    conf.addResource("tony-default.xml");
+    conf.setInt("tony.worker.instances", 1);
+    conf.setStrings("tony.containers.placement-spec", "java=true");
+    conf.setStrings("tony.worker.allocation-tags", "tony");
+    containerRequests = Utils.parseContainerRequests(conf);
+    assertEquals("java=true", containerRequests.get("worker").getPlacementSpec());
+    assertEquals(Arrays.asList("tony"), containerRequests.get("worker").getAllocationTags());
+  }
+
   @Test
   public void testIsArchive() {
     ClassLoader classLoader = getClass().getClassLoader();