Use allocationRequestId for matching returned resource requests to TF…

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

@@ -126,6 +126,10 @@ public class TonyApplicationMaster {
   private AtomicInteger numRequestedContainers = new AtomicInteger();
   private Map<String, List<ContainerRequest>> jobTypeToContainerRequestsMap = new HashMap<>();
 
+  // allocationRequestIds are allocated to TF tasks sequentially. lastAllocationRequestId
+  // tracks the latest allocationRequestId allocated.
+  private long lastAllocationRequestId = 0;
+
   // TensorFlow session
   private TensorFlowSession session = new TensorFlowSession(); // Create a dummy session for single node training.
   private TensorFlowSession.Builder sessionBuilder;
@@ -566,13 +570,13 @@ private boolean monitor() {
     }
     if (this.taskHasMissesHB) {
       session.setFinalStatus(FinalApplicationStatus.FAILED,
-                             "Application failed due to missed heartbeats");
+          "Application failed due to missed heartbeats");
     } else {
       session.updateSessionStatus();
     }
 
     LOG.info("Total completed worker tasks: " + numCompletedWorkerTasks.get()
-             + ", total worker tasks: " + numTotalWorkerTasks);
+        + ", total worker tasks: " + numTotalWorkerTasks);
     boolean success = true;
     FinalApplicationStatus status = session.getFinalStatus();
     String appMessage = session.getFinalMessage();
@@ -861,17 +865,14 @@ private void setupContainerCredentials() throws IOException {
   }
 
   private AMRMClient.ContainerRequest setupContainerRequestForRM(TensorFlowContainerRequest request) {
-    return setupContainerRequestForRM(request.getVirtualCores(), request.getMemory(), request.getGPU(),
-                                      request.getPriority());
-  }
-
-  private AMRMClient.ContainerRequest setupContainerRequestForRM(int vCores, long mem, final int gpu, int priority) {
-    Priority pi = Priority.newInstance(priority);
-    Resource capability = Resource.newInstance(mem, vCores);
-    Utils.setCapabilityGPU(capability, gpu);
-    AMRMClient.ContainerRequest request = new AMRMClient.ContainerRequest(capability, null, null, pi);
-    LOG.info("Requested container ask: " + request.toString());
-    return request;
+    Priority priority = Priority.newInstance(request.getPriority());
+    Resource capability = Resource.newInstance(request.getMemory(), request.getVirtualCores());
+    Utils.setCapabilityGPU(capability, request.getGPU());
+    session.addAllocationId(request.getJobName(), lastAllocationRequestId);
+    AMRMClient.ContainerRequest containerRequest = new AMRMClient.ContainerRequest(capability, null, null, priority,
+        lastAllocationRequestId++);
+    LOG.info("Requested container ask: " + containerRequest.toString());
+    return containerRequest;
   }
 
   private NMCallbackHandler createNMCallbackHandler() {
@@ -1022,7 +1023,7 @@ public void run() {
       containerEnv.put(Constants.SESSION_ID, String.valueOf(session.sessionId));
       Map<String, String> containerShellEnv = new ConcurrentHashMap<>(containerEnv);
 
-      TFTask task = session.getRemainingTask(container.getResource());
+      TFTask task = session.getMatchingTask(container.getAllocationRequestId());
 
       Preconditions.checkNotNull(task, "Task was null! Nothing to schedule.");
       task.addContainer(container);

tony-core/src/main/java/com/linkedin/tony/tensorflow/TensorFlowContainerRequest.java

@@ -51,20 +51,4 @@ public int getPriority() {
   public String getJobName() {
     return jobName;
   }
-
-  public boolean matchesResourceRequest(Resource resource) {
-    int resourceMem = (int) resource.getMemorySize();
-    int resourceVCore = resource.getVirtualCores();
-
-    if (this.memory != resourceMem || this.virtualCores != resourceVCore) {
-      return false;
-    }
-
-    try {
-      long numGpus = resource.getResourceValue(ResourceInformation.GPU_URI);
-      return this.gpu == numGpus;
-    } catch (ResourceNotFoundException e) {
-      return true;
-    }
-  }
 }

tony-core/src/main/java/com/linkedin/tony/tensorflow/TensorFlowSession.java

@@ -12,9 +12,11 @@
 import java.net.URI;
 import java.util.ArrayList;
 import java.util.HashMap;
+import java.util.HashSet;
 import java.util.List;
 import java.util.Map;
 import java.util.Objects;
+import java.util.Set;
 import java.util.concurrent.ConcurrentHashMap;
 import org.apache.commons.logging.Log;
 import org.apache.commons.logging.LogFactory;
@@ -62,7 +64,7 @@ public class TensorFlowSession {
   private Map<String, String> shellEnv;
   private String jvmArgs;
   private TensorFlowContainerRequest workerContainerRequest;
-  private Map<String, TensorFlowContainerRequest> jobTypeToRequestMap;
+  private Map<String, Set<Long>> jobTypeToAllocationIds = new HashMap<String, Set<Long>>();
 
   public enum TaskType {
     TASK_TYPE_CHIEF, TASK_TYPE_PARAMETER_SERVER, TASK_TYPE_OTHERS
@@ -112,8 +114,6 @@ private TensorFlowSession(Builder builder) {
 
     TFTask[] psTasks = new TFTask[this.numPs];
     jobTasks.put(PS_JOB_NAME, psTasks);
-
-    this.jobTypeToRequestMap = ImmutableMap.of("ps", psContainerRequest, "worker", workerContainerRequest);
   }
 
   public Map<String, TFTask[]> getTFTasks() {
@@ -211,17 +211,35 @@ public boolean allTasksScheduled() {
     return true;
   }
 
-  // Get a task that hasn't been scheduled.
-  public synchronized TFTask getRemainingTask(Resource resource) {
+  /**
+   * Associate an allocationRequestId to a TensorFlow job.
+   * @param jobName the TensorFlow job name
+   * @param allocationRequestId the allocationRequestId which corresponds to an instance of this job
+   */
+  public void addAllocationId(String jobName, long allocationRequestId) {
+    if (jobTypeToAllocationIds.get(jobName) == null) {
+      jobTypeToAllocationIds.put(jobName, new HashSet<>());
+    }
+    jobTypeToAllocationIds.get(jobName).add(allocationRequestId);
+    LOG.info(String.format("Job %s with allocationRequestId %d", jobName, allocationRequestId));
+  }
+
+  /**
+   * Get a TensorFlow task that hasn't been scheduled.
+   * @param allocationRequestId the allocationRequestId of the allocated container
+   * @return task to be assigned to this allocation
+   */
+  public synchronized TFTask getMatchingTask(long allocationRequestId) {
     for (Map.Entry<String, TFTask[]> entry : jobTasks.entrySet()) {
       String jobName = entry.getKey();
-      if (!jobTypeToRequestMap.get(jobName).matchesResourceRequest(resource)) {
+      if (!jobTypeToAllocationIds.get(jobName).contains(allocationRequestId)) {
         continue;
       }
       TFTask[] tasks = entry.getValue();
       for (int i = 0; i < tasks.length; i++) {
         if (tasks[i] == null) {
           tasks[i] = new TFTask(jobName, String.valueOf(i));
+          LOG.info(String.format("Matched job %s with allocationRequestId %d", jobName, allocationRequestId));
           return tasks[i];
         }
       }