org.apache.hadoop.yarn.api.records.ReservationId Java Examples

@Test
public void testSteps() {
  ReservationId reservationID =
      ReservationId.newInstance(rand.nextLong(), rand.nextLong());
  int[] alloc = { 10, 10, 10, 10, 10, 10 };
  int start = 100;
  ReservationDefinition rDef =
      createSimpleReservationDefinition(start, start + alloc.length + 1,
          alloc.length);
  Map<ReservationInterval, ReservationRequest> allocations =
      generateAllocation(start, alloc, true, false);
  ReservationAllocation rAllocation =
      new InMemoryReservationAllocation(reservationID, rDef, user, planName,
          start, start + alloc.length + 1, allocations, resCalc, minAlloc);
  doAssertions(rAllocation, reservationID, rDef, allocations, start, alloc);
  Assert.assertFalse(rAllocation.containsGangs());
  for (int i = 0; i < alloc.length; i++) {
    Assert.assertEquals(
        Resource.newInstance(1024 * (alloc[i] + i), (alloc[i] + i)),
        rAllocation.getResourcesAtTime(start + i));
  }
}
 

@Test
public void testAddEmptyReservation() {
  Plan plan =
      new InMemoryPlan(queueMetrics, policy, agent, totalCapacity, 1L,
          resCalc, minAlloc, maxAlloc, planName, replanner, true);
  ReservationId reservationID =
      ReservationSystemTestUtil.getNewReservationId();
  int[] alloc = {};
  int start = 100;
  Map<ReservationInterval, ReservationRequest> allocations =
      new HashMap<ReservationInterval, ReservationRequest>();
  ReservationDefinition rDef =
      createSimpleReservationDefinition(start, start + alloc.length,
          alloc.length, allocations.values());
  ReservationAllocation rAllocation =
      new InMemoryReservationAllocation(reservationID, rDef, user, planName,
          start, start + alloc.length, allocations, resCalc, minAlloc);
  Assert.assertNull(plan.getReservationById(reservationID));
  try {
    plan.addReservation(rAllocation);
  } catch (PlanningException e) {
    Assert.fail(e.getMessage());
  }
}
 

@Override
public boolean deleteReservation(ReservationId reservationID) {
  writeLock.lock();
  try {
    ReservationAllocation reservation = getReservationById(reservationID);
    if (reservation == null) {
      String errMsg =
          "The specified Reservation with ID " + reservationID
              + " does not exist in the plan";
      LOG.error(errMsg);
      throw new IllegalArgumentException(errMsg);
    }
    return removeReservation(reservation);
  } finally {
    writeLock.unlock();
  }
}
 

@Test
public void testAddEmptyReservation() {
  Plan plan =
      new InMemoryPlan(queueMetrics, policy, agent, totalCapacity, 1L,
          resCalc, minAlloc, maxAlloc, planName, replanner, true);
  ReservationId reservationID =
      ReservationSystemTestUtil.getNewReservationId();
  int[] alloc = {};
  int start = 100;
  Map<ReservationInterval, ReservationRequest> allocations =
      new HashMap<ReservationInterval, ReservationRequest>();
  ReservationDefinition rDef =
      createSimpleReservationDefinition(start, start + alloc.length,
          alloc.length, allocations.values());
  ReservationAllocation rAllocation =
      new InMemoryReservationAllocation(reservationID, rDef, user, planName,
          start, start + alloc.length, allocations, resCalc, minAlloc);
  Assert.assertNull(plan.getReservationById(reservationID));
  try {
    plan.addReservation(rAllocation);
  } catch (PlanningException e) {
    Assert.fail(e.getMessage());
  }
}
 

@Test
public void testSkyline() {
  ReservationId reservationID =
      ReservationId.newInstance(rand.nextLong(), rand.nextLong());
  int[] alloc = { 0, 5, 10, 10, 5, 0 };
  int start = 100;
  ReservationDefinition rDef =
      createSimpleReservationDefinition(start, start + alloc.length + 1,
          alloc.length);
  Map<ReservationInterval, ReservationRequest> allocations =
      generateAllocation(start, alloc, true, false);
  ReservationAllocation rAllocation =
      new InMemoryReservationAllocation(reservationID, rDef, user, planName,
          start, start + alloc.length + 1, allocations, resCalc, minAlloc);
  doAssertions(rAllocation, reservationID, rDef, allocations, start, alloc);
  Assert.assertFalse(rAllocation.containsGangs());
  for (int i = 0; i < alloc.length; i++) {
    Assert.assertEquals(
        Resource.newInstance(1024 * (alloc[i] + i), (alloc[i] + i)),
        rAllocation.getResourcesAtTime(start + i));
  }
}
 

InMemoryReservationAllocation(ReservationId reservationID,
    ReservationDefinition contract, String user, String planName,
    long startTime, long endTime,
    Map<ReservationInterval, ReservationRequest> allocationRequests,
    ResourceCalculator calculator, Resource minAlloc) {
  this.contract = contract;
  this.startTime = startTime;
  this.endTime = endTime;
  this.reservationID = reservationID;
  this.user = user;
  this.allocationRequests = allocationRequests;
  this.planName = planName;
  resourcesOverTime = new RLESparseResourceAllocation(calculator, minAlloc);
  for (Map.Entry<ReservationInterval, ReservationRequest> r : allocationRequests
      .entrySet()) {
    resourcesOverTime.addInterval(r.getKey(), r.getValue());
    if (r.getValue().getConcurrency() > 1) {
      hasGang = true;
    }
  }
}
 

@Test
public void testUpdateReservationDoesnotExist() {
  ReservationUpdateRequest request =
      createSimpleReservationUpdateRequest(1, 1, 1, 5, 4);
  ReservationId rId = request.getReservationId();
  when(rSystem.getQueueForReservation(rId)).thenReturn(null);
  Plan plan = null;
  try {
    plan = rrValidator.validateReservationUpdateRequest(rSystem, request);
    Assert.fail();
  } catch (YarnException e) {
    Assert.assertNull(plan);
    String message = e.getMessage();
    Assert
        .assertTrue(message.equals(MessageFormat
            .format(
                "The specified reservation with ID: {0} is unknown. Please try again with a valid reservation.",
                rId)));
    LOG.info(message);
  }
}
 

@Test
public void testDeleteReservationDoesnotExist() {
  ReservationDeleteRequest request = new ReservationDeleteRequestPBImpl();
  ReservationId rId = ReservationSystemTestUtil.getNewReservationId();
  request.setReservationId(rId);
  when(rSystem.getQueueForReservation(rId)).thenReturn(null);
  Plan plan = null;
  try {
    plan = rrValidator.validateReservationDeleteRequest(rSystem, request);
    Assert.fail();
  } catch (YarnException e) {
    Assert.assertNull(plan);
    String message = e.getMessage();
    Assert
        .assertTrue(message.equals(MessageFormat
            .format(
                "The specified reservation with ID: {0} is unknown. Please try again with a valid reservation.",
                rId)));
    LOG.info(message);
  }
}
 

@Test
public void testDeleteReservationDoesnotExist() {
  ReservationDeleteRequest request = new ReservationDeleteRequestPBImpl();
  ReservationId rId = ReservationSystemTestUtil.getNewReservationId();
  request.setReservationId(rId);
  when(rSystem.getQueueForReservation(rId)).thenReturn(null);
  Plan plan = null;
  try {
    plan = rrValidator.validateReservationDeleteRequest(rSystem, request);
    Assert.fail();
  } catch (YarnException e) {
    Assert.assertNull(plan);
    String message = e.getMessage();
    Assert
        .assertTrue(message.equals(MessageFormat
            .format(
                "The specified reservation with ID: {0} is unknown. Please try again with a valid reservation.",
                rId)));
    LOG.info(message);
  }
}
 

@Test
public void testDeleteReservationInvalidPlan() {
  ReservationDeleteRequest request = new ReservationDeleteRequestPBImpl();
  ReservationId reservationID =
      ReservationSystemTestUtil.getNewReservationId();
  request.setReservationId(reservationID);
  when(rSystem.getPlan(PLAN_NAME)).thenReturn(null);
  Plan plan = null;
  try {
    plan = rrValidator.validateReservationDeleteRequest(rSystem, request);
    Assert.fail();
  } catch (YarnException e) {
    Assert.assertNull(plan);
    String message = e.getMessage();
    Assert
        .assertTrue(message
            .endsWith(" is not associated with any valid plan. Please try again with a valid reservation."));
    LOG.info(message);
  }
}
 

@Test
public void testGangAlloaction() {
  ReservationId reservationID =
      ReservationId.newInstance(rand.nextLong(), rand.nextLong());
  int[] alloc = { 10, 10, 10, 10, 10, 10 };
  int start = 100;
  ReservationDefinition rDef =
      createSimpleReservationDefinition(start, start + alloc.length + 1,
          alloc.length);
  Map<ReservationInterval, ReservationRequest> allocations =
      generateAllocation(start, alloc, false, true);
  ReservationAllocation rAllocation =
      new InMemoryReservationAllocation(reservationID, rDef, user, planName,
          start, start + alloc.length + 1, allocations, resCalc, minAlloc);
  doAssertions(rAllocation, reservationID, rDef, allocations, start, alloc);
  Assert.assertTrue(rAllocation.containsGangs());
  for (int i = 0; i < alloc.length; i++) {
    Assert.assertEquals(Resource.newInstance(1024 * (alloc[i]), (alloc[i]), (alloc[i])),
        rAllocation.getResourcesAtTime(start + i));
  }
}
 

@Test
public void testZeroAlloaction() {
  ReservationId reservationID =
      ReservationId.newInstance(rand.nextLong(), rand.nextLong());
  int[] alloc = {};
  long start = 0;
  ReservationDefinition rDef =
      createSimpleReservationDefinition(start, start + alloc.length + 1,
          alloc.length);
  Map<ReservationInterval, ReservationRequest> allocations =
      new HashMap<ReservationInterval, ReservationRequest>();
  ReservationAllocation rAllocation =
      new InMemoryReservationAllocation(reservationID, rDef, user, planName,
          start, start + alloc.length + 1, allocations, resCalc, minAlloc);
  doAssertions(rAllocation, reservationID, rDef, allocations, (int) start,
      alloc);
  Assert.assertFalse(rAllocation.containsGangs());
}
 

@Before
public void setUp() {
  clock = mock(Clock.class);
  plan = mock(Plan.class);
  rSystem = mock(ReservationSystem.class);
  plans.put(PLAN_NAME, plan);
  rrValidator = new ReservationInputValidator(clock);
  when(clock.getTime()).thenReturn(1L);
  ResourceCalculator rCalc = new DefaultResourceCalculator();
  Resource resource = Resource.newInstance(10240, 10);
  when(plan.getResourceCalculator()).thenReturn(rCalc);
  when(plan.getTotalCapacity()).thenReturn(resource);
  when(rSystem.getQueueForReservation(any(ReservationId.class))).thenReturn(
      PLAN_NAME);
  when(rSystem.getPlan(PLAN_NAME)).thenReturn(plan);
}
 

@Test
public void testSteps() {
  ReservationId reservationID =
      ReservationId.newInstance(rand.nextLong(), rand.nextLong());
  int[] alloc = { 10, 10, 10, 10, 10, 10 };
  int start = 100;
  ReservationDefinition rDef =
      createSimpleReservationDefinition(start, start + alloc.length + 1,
          alloc.length);
  Map<ReservationInterval, ReservationRequest> allocations =
      generateAllocation(start, alloc, true, false);
  ReservationAllocation rAllocation =
      new InMemoryReservationAllocation(reservationID, rDef, user, planName,
          start, start + alloc.length + 1, allocations, resCalc, minAlloc);
  doAssertions(rAllocation, reservationID, rDef, allocations, start, alloc);
  Assert.assertFalse(rAllocation.containsGangs());
  for (int i = 0; i < alloc.length; i++) {
    Assert.assertEquals(
        Resource.newInstance(1024 * (alloc[i] + i), (alloc[i] + i), (alloc[i] + i)),
        rAllocation.getResourcesAtTime(start + i));
  }
}
 

@Test(expected = MismatchedUserException.class)
public void testUserMismatch() throws IOException, PlanningException {
  // generate allocation from single tenant that exceed capacity
  int[] f = generateData(3600, (int) (0.5 * totCont));

  ReservationId rid = ReservationSystemTestUtil.getNewReservationId();
  plan.addReservation(new InMemoryReservationAllocation(rid, null, "u1",
      "dedicated", initTime, initTime + f.length, ReservationSystemTestUtil
          .generateAllocation(initTime, step, f), res, minAlloc));

  // trying to update a reservation with a mismatching user
  plan.updateReservation(new InMemoryReservationAllocation(rid, null, "u2",
      "dedicated", initTime, initTime + f.length, ReservationSystemTestUtil
          .generateAllocation(initTime, step, f), res, minAlloc));
}
 

@Override
protected Resource getReservationQueueResourceIfExists(Plan plan,
    ReservationId reservationId) {
  CSQueue resQueue = cs.getQueue(reservationId.toString());
  Resource reservationResource = null;
  if (resQueue != null) {
    reservationResource = Resources.multiply(cs.getClusterResource(),
        resQueue.getAbsoluteCapacity());
  }
  return reservationResource;
}
 

@Override
public ReservationId getReservationID() {
  ApplicationSubmissionContextProtoOrBuilder p = viaProto ? proto : builder;
  if (reservationId != null) {
    return reservationId;
  }
  if (!p.hasReservationId()) {
    return null;
  }
  reservationId = convertFromProtoFormat(p.getReservationId());
  return reservationId;
}
 

@Override
public ReservationDeleteResponse deleteReservation(
    ReservationDeleteRequest request) throws YarnException, IOException {
  // Check if reservation system is enabled
  checkReservationSytem(AuditConstants.DELETE_RESERVATION_REQUEST);
  ReservationDeleteResponse response =
      recordFactory.newRecordInstance(ReservationDeleteResponse.class);
  // Validate the input
  Plan plan =
      rValidator.validateReservationDeleteRequest(reservationSystem, request);
  ReservationId reservationId = request.getReservationId();
  String queueName = reservationSystem.getQueueForReservation(reservationId);
  // Check ACLs
  String user =
      checkReservationACLs(queueName,
          AuditConstants.DELETE_RESERVATION_REQUEST);
  // Try to update the reservation using default agent
  try {
    boolean result =
        plan.getReservationAgent().deleteReservation(reservationId, user,
            plan);
    if (!result) {
      String errMsg = "Could not delete reservation: " + reservationId;
      RMAuditLogger.logFailure(user,
          AuditConstants.DELETE_RESERVATION_REQUEST, errMsg,
          "ClientRMService", errMsg);
      throw RPCUtil.getRemoteException(errMsg);
    }
  } catch (PlanningException e) {
    RMAuditLogger.logFailure(user, AuditConstants.DELETE_RESERVATION_REQUEST,
        e.getMessage(), "ClientRMService",
        "Unable to delete the reservation: " + reservationId);
    throw RPCUtil.getRemoteException(e);
  }
  RMAuditLogger.logSuccess(user, AuditConstants.DELETE_RESERVATION_REQUEST,
      "ClientRMService: " + reservationId);
  return response;
}
 

@Override
public void setReservationId(ReservationId reservationId) {
  maybeInitBuilder();
  if (reservationId == null) {
    builder.clearReservationId();
    return;
  }
  this.reservationId = reservationId;
}
 

private Plan validateReservation(ReservationSystem reservationSystem,
    ReservationId reservationId, String auditConstant) throws YarnException {
  String message = "";
  // check if the reservation id is valid
  if (reservationId == null) {
    message =
        "Missing reservation id."
            + " Please try again by specifying a reservation id.";
    RMAuditLogger.logFailure("UNKNOWN", auditConstant,
        "validate reservation input", "ClientRMService", message);
    throw RPCUtil.getRemoteException(message);
  }
  String queueName = reservationSystem.getQueueForReservation(reservationId);
  if (queueName == null) {
    message =
        "The specified reservation with ID: " + reservationId
            + " is unknown. Please try again with a valid reservation.";
    RMAuditLogger.logFailure("UNKNOWN", auditConstant,
        "validate reservation input", "ClientRMService", message);
    throw RPCUtil.getRemoteException(message);
  }
  // check if the associated plan is valid
  Plan plan = reservationSystem.getPlan(queueName);
  if (plan == null) {
    message =
        "The specified reservation: " + reservationId
            + " is not associated with any valid plan."
            + " Please try again with a valid reservation.";
    RMAuditLogger.logFailure("UNKNOWN", auditConstant,
        "validate reservation input", "ClientRMService", message);
    throw RPCUtil.getRemoteException(message);
  }
  return plan;
}
 

@Override
public void setReservationId(ReservationId reservationId) {
  maybeInitBuilder();
  if (reservationId == null) {
    builder.clearReservationId();
    return;
  }
  this.reservationId = reservationId;
}
 

private void doAssertions(Plan plan, ReservationAllocation rAllocation) {
  ReservationId reservationID = rAllocation.getReservationId();
  Assert.assertNotNull(plan.getReservationById(reservationID));
  Assert.assertEquals(rAllocation, plan.getReservationById(reservationID));
  Assert.assertTrue(((InMemoryPlan) plan).getAllReservations().size() == 1);
  Assert.assertEquals(rAllocation.getEndTime(), plan.getLastEndTime());
  Assert.assertEquals(totalCapacity, plan.getTotalCapacity());
  Assert.assertEquals(minAlloc, plan.getMinimumAllocation());
  Assert.assertEquals(maxAlloc, plan.getMaximumAllocation());
  Assert.assertEquals(resCalc, plan.getResourceCalculator());
  Assert.assertEquals(planName, plan.getQueueName());
  Assert.assertTrue(plan.getMoveOnExpiry());
}
 

@Public
@Unstable
public static ReservationDeleteRequest newInstance(ReservationId reservationId) {
  ReservationDeleteRequest request =
      Records.newRecord(ReservationDeleteRequest.class);
  request.setReservationId(reservationId);
  return request;
}
 

@Override
public ReservationId getReservationId() {
  ReservationDeleteRequestProtoOrBuilder p = viaProto ? proto : builder;
  if (reservationId != null) {
    return reservationId;
  }
  if (!p.hasReservationId()) {
    return null;
  }
  reservationId = convertFromProtoFormat(p.getReservationId());
  return reservationId;
}
 

@Override
public String getQueueForReservation(ReservationId reservationId) {
  readLock.lock();
  try {
    return resQMap.get(reservationId);
  } finally {
    readLock.unlock();
  }
}
 

@Test
public void testSingleSliding() throws PlanningException {
  prepareBasicPlan();

  // create a single request for which we need subsequent (tight) packing.
  ReservationDefinition rr = new ReservationDefinitionPBImpl();
  rr.setArrival(100 * step);
  rr.setDeadline(120 * step);
  ReservationRequests reqs = new ReservationRequestsPBImpl();
  reqs.setInterpreter(ReservationRequestInterpreter.R_ALL);
  ReservationRequest r = ReservationRequest.newInstance(
      Resource.newInstance(1024, 1, 1), 200, 10, 10 * step);

  List<ReservationRequest> list = new ArrayList<ReservationRequest>();
  list.add(r);
  reqs.setReservationResources(list);
  rr.setReservationRequests(reqs);

  // submit to agent
  ReservationId reservationID = ReservationSystemTestUtil
      .getNewReservationId();
  agent.createReservation(reservationID, "u1", plan, rr);

  // validate results, we expect the second one to be accepted
  assertTrue("Agent-based allocation failed", reservationID != null);
  assertTrue("Agent-based allocation failed", plan.getAllReservations()
      .size() == 3);

  ReservationAllocation cs = plan.getReservationById(reservationID);

  assertTrue(cs.toString(), check(cs, 100 * step, 120 * step, 100, 1024, 1));

  System.out.println("--------AFTER packed ALLOCATION (queue: "
      + reservationID + ")----------");
  System.out.println(plan.toString());
  System.out.println(plan.toCumulativeString());

}
 

@Override
public ReservationId getNewReservationId() {
  writeLock.lock();
  try {
    ReservationId resId =
        ReservationId.newInstance(ResourceManager.getClusterTimeStamp(),
            resCounter.incrementAndGet());
    LOG.info("Allocated new reservationId: " + resId);
    return resId;
  } finally {
    writeLock.unlock();
  }
}
 

@Override
public ReservationId getReservationId() {
  ReservationUpdateRequestProtoOrBuilder p = viaProto ? proto : builder;
  if (reservationId != null) {
    return reservationId;
  }
  if (!p.hasReservationId()) {
    return null;
  }
  reservationId = convertFromProtoFormat(p.getReservationId());
  return reservationId;
}
 

@Override
public ReservationUpdateResponse updateReservation(
    ReservationUpdateRequest request) throws YarnException, IOException {
  // Check if reservation system is enabled
  checkReservationSytem(AuditConstants.UPDATE_RESERVATION_REQUEST);
  ReservationUpdateResponse response =
      recordFactory.newRecordInstance(ReservationUpdateResponse.class);
  // Validate the input
  Plan plan =
      rValidator.validateReservationUpdateRequest(reservationSystem, request);
  ReservationId reservationId = request.getReservationId();
  String queueName = reservationSystem.getQueueForReservation(reservationId);
  // Check ACLs
  String user =
      checkReservationACLs(queueName,
          AuditConstants.UPDATE_RESERVATION_REQUEST);
  // Try to update the reservation using default agent
  try {
    boolean result =
        plan.getReservationAgent().updateReservation(reservationId, user,
            plan, request.getReservationDefinition());
    if (!result) {
      String errMsg = "Unable to update reservation: " + reservationId;
      RMAuditLogger.logFailure(user,
          AuditConstants.UPDATE_RESERVATION_REQUEST, errMsg,
          "ClientRMService", errMsg);
      throw RPCUtil.getRemoteException(errMsg);
    }
  } catch (PlanningException e) {
    RMAuditLogger.logFailure(user, AuditConstants.UPDATE_RESERVATION_REQUEST,
        e.getMessage(), "ClientRMService",
        "Unable to update the reservation: " + reservationId);
    throw RPCUtil.getRemoteException(e);
  }
  RMAuditLogger.logSuccess(user, AuditConstants.UPDATE_RESERVATION_REQUEST,
      "ClientRMService: " + reservationId);
  return response;
}
 

@Override
public String getQueueForReservation(ReservationId reservationId) {
  readLock.lock();
  try {
    return resQMap.get(reservationId);
  } finally {
    readLock.unlock();
  }
}