Carry forward resource overrides when retrying run-nows.

SingularityService/src/main/java/com/hubspot/singularity/scheduler/SingularityScheduler.java

@@ -36,6 +36,7 @@
 import com.google.common.collect.Sets;
 import com.google.inject.Inject;
 import com.hubspot.mesos.JavaUtils;
+import com.hubspot.mesos.Resources;
 import com.hubspot.singularity.DeployState;
 import com.hubspot.singularity.ExtendedTaskState;
 import com.hubspot.singularity.MachineState;
@@ -590,12 +591,14 @@ private Optional<PendingType> handleCompletedTaskWithStatistics(Optional<Singula
 
     PendingType pendingType = PendingType.TASK_DONE;
     Optional<List<String>> cmdLineArgsList = Optional.absent();
+    Optional<Resources> resources = Optional.absent();
 
     if (!state.isSuccess() && shouldRetryImmediately(request, deployStatistics, task)) {
       LOG.debug("Retrying {} because {}", request.getId(), state);
       pendingType = PendingType.RETRY;
       if (task.isPresent()) {
         cmdLineArgsList = task.get().getTaskRequest().getPendingTask().getCmdLineArgsList();
+        resources = task.get().getTaskRequest().getPendingTask().getResources();
       }
     } else if (!request.isAlwaysRunning()) {
       return Optional.absent();
@@ -609,7 +612,7 @@ private Optional<PendingType> handleCompletedTaskWithStatistics(Optional<Singula
 
     SingularityPendingRequest pendingRequest = new SingularityPendingRequest(request.getId(), requestDeployState.get().getActiveDeploy().get().getDeployId(),
         System.currentTimeMillis(), Optional.absent(), pendingType, cmdLineArgsList, Optional.absent(), Optional.absent(), Optional.absent(),
-        Optional.absent());
+        Optional.absent(), resources, Collections.emptyList(), Optional.absent(), Collections.emptyMap(), Collections.emptyList(), Optional.absent());
 
     requestManager.addToPendingQueue(pendingRequest);
 

SingularityService/src/test/java/com/hubspot/singularity/scheduler/SingularitySchedulerTest.java

@@ -758,6 +758,65 @@ public void testRetries() {
     Assert.assertTrue(taskManager.getActiveTaskIds().isEmpty());
   }
 
+  @Test
+  public void testRetriesWithOverrides() {
+    SingularityRequestBuilder bldr = new SingularityRequestBuilder(requestId, RequestType.ON_DEMAND);
+    request = bldr.setNumRetriesOnFailure(Optional.of(2)).build();
+    saveRequest(request);
+
+    deployResource.deploy(new SingularityDeployRequest(new SingularityDeployBuilder(requestId, "d1").setCommand(Optional.of("cmd")).build(), Optional.absent(), Optional.absent()), singularityUser);
+
+    scheduler.drainPendingQueue();
+    deployChecker.checkDeploys();
+
+    requestResource.scheduleImmediately(
+        singularityUser,
+        requestId,
+        new SingularityRunNowRequestBuilder()
+            .setCommandLineArgs(Collections.singletonList("extraFlag"))
+            .setResources(new Resources(17, 1337, 0))
+            .build()
+    );
+
+    resourceOffers();
+
+    Assert.assertEquals(1, taskManager.getActiveTaskIds().size());
+
+    Resources resourcesForRunningTask = taskManager.getActiveTasks().get(0).getTaskRequest().getPendingTask().getResources().get();
+    Assert.assertEquals(Optional.of(Collections.singletonList("extraFlag")), taskManager.getActiveTasks().get(0).getTaskRequest().getPendingTask().getCmdLineArgsList());
+    Assert.assertEquals(17, resourcesForRunningTask.getCpus(), 0.01);
+    Assert.assertEquals(1337, resourcesForRunningTask.getMemoryMb(), 0.01);
+
+    statusUpdate(taskManager.getActiveTasks().get(0), TaskState.TASK_LOST);
+
+    resourceOffers();
+
+    Assert.assertEquals(1, taskManager.getActiveTaskIds().size());
+
+    resourcesForRunningTask = taskManager.getActiveTasks().get(0).getTaskRequest().getPendingTask().getResources().get();
+    Assert.assertEquals(Optional.of(Collections.singletonList("extraFlag")), taskManager.getActiveTasks().get(0).getTaskRequest().getPendingTask().getCmdLineArgsList());
+    Assert.assertEquals(17, resourcesForRunningTask.getCpus(), 0.01);
+    Assert.assertEquals(1337, resourcesForRunningTask.getMemoryMb(), 0.01);
+
+    statusUpdate(taskManager.getActiveTasks().get(0), TaskState.TASK_LOST);
+
+    resourceOffers();
+
+    Assert.assertEquals(1, taskManager.getActiveTaskIds().size());
+
+    resourcesForRunningTask = taskManager.getActiveTasks().get(0).getTaskRequest().getPendingTask().getResources().get();
+    Assert.assertEquals(Optional.of(Collections.singletonList("extraFlag")), taskManager.getActiveTasks().get(0).getTaskRequest().getPendingTask().getCmdLineArgsList());
+    Assert.assertEquals(17, resourcesForRunningTask.getCpus(), 0.01);
+    Assert.assertEquals(1337, resourcesForRunningTask.getMemoryMb(), 0.01);
+
+    statusUpdate(taskManager.getActiveTasks().get(0), TaskState.TASK_LOST);
+
+    resourceOffers();
+
+    Assert.assertTrue(taskManager.getActiveTaskIds().isEmpty());
+
+  }
+
   @Test
   public void testCooldownAfterSequentialFailures() {
     initRequest();