public GpuAllocation assignGpus(Container container)
      throws ResourceHandlerException {
    GpuAllocation allocation = internalAssignGpus(container);

    // Wait for a maximum of waitPeriodForResource seconds if no
    // available GPU are there which are yet to be released.
    int timeWaiting = 0;
    while (allocation == null) {
      if (timeWaiting >= waitPeriodForResource) {
        break;
      }

      // Sleep for 1 sec to ensure there are some free GPU devices which are
      // getting released.
      try {
        LOG.info("Container : " + container.getContainerId()
            + " is waiting for free GPU devices.");
        Thread.sleep(WAIT_MS_PER_LOOP);
        timeWaiting += WAIT_MS_PER_LOOP;
        allocation = internalAssignGpus(container);
      } catch (InterruptedException e) {
        // On any interrupt, break the loop and continue execution.
        Thread.currentThread().interrupt();
        LOG.warn("Interrupted while waiting for available GPU");
        break;
      }
    }

    if(allocation == null) {
      String message = "Could not get valid GPU device for container '" +
          container.getContainerId()
          + "' as some other containers might not releasing GPUs.";
      LOG.warn(message);
      throw new ResourceHandlerException(message);
    }
    return allocation;
  }
@Test
  public void testGpuGetsUnassignedWhenStateStoreThrowsException()
      throws ResourceHandlerException, IOException {
    doThrow(new IOException("Failed to save container mappings " +
        "to NM state store!"))
        .when(nmStateStore).storeAssignedResources(any(Container.class),
        anyString(), anyList());

    createAndAddGpus(1);

    exception.expect(ResourceHandlerException.class);
    exception.expectMessage("Failed to save container mappings " +
        "to NM state store");
    Container container = createMockContainer(1, 5L);
    testSubject.assignGpus(container);
  }