Package edu.berkeley.sparrow.daemon.scheduler

Source Code of edu.berkeley.sparrow.daemon.scheduler.TestConstrainedTaskPlacer

/*
* Copyright 2013 The Regents of The University California
*
* Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License.
* You may obtain a copy of the License at
*
*   http://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and
* limitations under the License.
*/

package edu.berkeley.sparrow.daemon.scheduler;
import static org.junit.Assert.assertEquals;
import static org.junit.Assert.assertTrue;

import java.net.InetSocketAddress;
import java.nio.ByteBuffer;
import java.util.ArrayList;
import java.util.HashSet;
import java.util.List;
import java.util.Map;
import java.util.Map.Entry;
import java.util.Set;

import org.apache.log4j.BasicConfigurator;
import org.junit.After;
import org.junit.Before;
import org.junit.Test;

import edu.berkeley.sparrow.daemon.util.Network;
import edu.berkeley.sparrow.thrift.TEnqueueTaskReservationsRequest;
import edu.berkeley.sparrow.thrift.THostPort;
import edu.berkeley.sparrow.thrift.TPlacementPreference;
import edu.berkeley.sparrow.thrift.TSchedulingRequest;
import edu.berkeley.sparrow.thrift.TTaskLaunchSpec;
import edu.berkeley.sparrow.thrift.TTaskSpec;
import edu.berkeley.sparrow.thrift.TUserGroupInfo;

public class TestConstrainedTaskPlacer {
  private static final String APP_ID = "test app";
  private static final String USER = "user";
  private static final String GROUP = "group";
  private static final int PRIORITY = 0;
  private static final String REQUEST_ID = "request id";
  private static final THostPort SCHEDULER_ADDRESS = new THostPort("localhost", 12345);
  private static final int ITERATIONS = 1;

  private static final TUserGroupInfo user = new TUserGroupInfo(USER, GROUP, PRIORITY);
  private static final Set<InetSocketAddress> preferredNodes = new HashSet<InetSocketAddress>();
  private static final Set<InetSocketAddress> allBackends = new HashSet<InetSocketAddress>();

  @Before
  public void setUp() {
    // Set up a simple configuration that logs on the console.
    BasicConfigurator.configure();
  }

  @After
  public void tearDown() {
    preferredNodes.clear();
    allBackends.clear();
  }

  /** Sanity checks a list of requests returned by {@code getEnqueueTaskReservationsRequests()}.
   *
   * For each request, ensures that the reservation is on one of the preferred nodes, and that it
   * has the appropriate app id/user/request id/scheduler address.
   */
  private void sanityCheckRequests(
      Map<InetSocketAddress, TEnqueueTaskReservationsRequest> requests,
      int expectedTotalReservations) {
    int totalEnqueuedTasks = 0;
    for (Entry<InetSocketAddress, TEnqueueTaskReservationsRequest> entry : requests.entrySet()) {
      // The node monitor the request is being sent to should be among the list of preferred
      // nodes.
      assertTrue("Expect " + entry.getKey() + " to be among the set of preferred nodes",
                 preferredNodes.contains(entry.getKey()));
      TEnqueueTaskReservationsRequest request = entry.getValue();
      assertEquals(request.getAppId(), APP_ID);
      assertEquals(request.getUser(), user);
      assertEquals(request.getRequestId(), REQUEST_ID);
      assertEquals(request.getSchedulerAddress(), SCHEDULER_ADDRESS);
      assertTrue(request.getNumTasks() > 0);
      totalEnqueuedTasks += request.getNumTasks();
    }
    assertEquals(expectedTotalReservations, totalEnqueuedTasks);
  }

  /** Generates a scheduling request with 3 tasks, and populates preferredNodes/allBackends. */
  private TSchedulingRequest generateSchedulingRequests() {
    preferredNodes.add(new InetSocketAddress("123.4.5.6", 1));
    preferredNodes.add(new InetSocketAddress("3.4.5.6", 56));
    preferredNodes.add(new InetSocketAddress("4.2.3.67", 89));
    preferredNodes.add(new InetSocketAddress("4.1.23.45", 100));
    preferredNodes.add(new InetSocketAddress("9.0.0.0", 204));

    List<InetSocketAddress> preferredNodesList = new ArrayList<InetSocketAddress>(preferredNodes);

    // Create 3 tasks with overlapping placement preferences (all tasks include
    // preferredNodesList[2] in their placement preferences).
    ByteBuffer message = ByteBuffer.allocate(1);
    List<TTaskSpec> tasks = new ArrayList<TTaskSpec>();
    final int NUM_TASKS = 3;
    for (int i = 0; i < NUM_TASKS; ++i) {
      TPlacementPreference placementPreference = new TPlacementPreference();
      for (int j = i; j < i + 3; ++j) {
        placementPreference.addToNodes(preferredNodesList.get(j).getAddress().getHostAddress());
      }
      String id = "test task " + i;
      tasks.add(new TTaskSpec(id, placementPreference, message));
    }

    TSchedulingRequest schedulingRequest = new TSchedulingRequest(APP_ID, tasks, user);

    // Create list of available backend nodes, with some additional nodes in addition to the
    // preferred ones.
    allBackends.addAll(preferredNodes);
    allBackends.add(new InetSocketAddress("1.2.3.4", 345));
    allBackends.add(new InetSocketAddress("5.6.7.81", 2000));
    return schedulingRequest;
  }

  /**
   * Creates a scheduling request with a single task that has 3 preferred nodes, and ensures
   * that repeated calls to getEnqueueTaskReservationsRequests always returns requests for one
   * of those three nodes.
   */
  @Test
  public void testGetEnqueueTaskReservationsRequestsSingleTask() {
    ConstrainedTaskPlacer taskPlacer = new ConstrainedTaskPlacer(REQUEST_ID, 2);

    preferredNodes.add(new InetSocketAddress("127.0.0.1", 22));
    preferredNodes.add(new InetSocketAddress("123.4.5.6", 20000));
    preferredNodes.add(new InetSocketAddress("7.0.0.9", 45));

    // Create a single task with three placement preferences.
    ByteBuffer message = ByteBuffer.allocate(1);
    TPlacementPreference placementPreference = new TPlacementPreference();
    for (InetSocketAddress address : preferredNodes) {
      placementPreference.addToNodes(address.getAddress().getHostAddress());
    }
    TTaskSpec task = new TTaskSpec("test task", placementPreference, message);

    // Create the scheduling request.
    List<TTaskSpec> tasks = new ArrayList<TTaskSpec>();
    tasks.add(task);
    TSchedulingRequest schedulingRequest = new TSchedulingRequest(APP_ID, tasks, user);

    // Create a set of backends, including all of the preferred nodes and some extra nodes too.
    allBackends.add(new InetSocketAddress("3.4.5.6", 174));
    for (InetSocketAddress node : preferredNodes) {
      allBackends.add(node);
    }
    allBackends.add(new InetSocketAddress("234.5.6.7", 22));
    allBackends.add(new InetSocketAddress("9.8.7.6", 1));

    for (int i = 0; i < ITERATIONS; ++i) {
      Map<InetSocketAddress, TEnqueueTaskReservationsRequest> requests =
          taskPlacer.getEnqueueTaskReservationsRequests(schedulingRequest, REQUEST_ID, allBackends,
                                                        SCHEDULER_ADDRESS);
      sanityCheckRequests(requests, 2);
    }
  }

  /**
   * Creates a scheduling request with 3 tasks, and sanity checks the result of
   * getEnqueueTaskReservationsRequests()
   */
  @Test
  public void testMultipleTasks() {
    TSchedulingRequest schedulingRequest = generateSchedulingRequests();

    for (int i = 0; i < ITERATIONS; ++i) {
      ConstrainedTaskPlacer taskPlacer = new ConstrainedTaskPlacer(REQUEST_ID, 2);

      Map<InetSocketAddress, TEnqueueTaskReservationsRequest> requests =
          taskPlacer.getEnqueueTaskReservationsRequests(schedulingRequest, REQUEST_ID,
                                                        allBackends, SCHEDULER_ADDRESS);
      sanityCheckRequests(requests, 6);
      sanityCheckGetTasksAndCancellation(schedulingRequest, taskPlacer,
          requests);
    }

  }

  private void sanityCheckGetTasksAndCancellation(
      TSchedulingRequest schedulingRequest, ConstrainedTaskPlacer taskPlacer,
      Map<InetSocketAddress, TEnqueueTaskReservationsRequest> requests) {
    Set<String> taskIds = new HashSet<String>();
    for (TTaskSpec task : schedulingRequest.tasks) {
      taskIds.add(task.taskId);
    }

    // Now try to get the three tasks.
    Set<String> unlaunchedTaskIds = new HashSet<String>(taskIds);
    int numGetTasks = 0;
    Set<THostPort> cancellations = null;
    for (Entry<InetSocketAddress, TEnqueueTaskReservationsRequest> entry : requests.entrySet()) {
      // For each entry, try to assign the tasks.  Once all tasks have been assigned, cancel the
      // remaining reservations.

      THostPort hostPort = Network.socketAddressToThrift(entry.getKey());
      // If the request has been cancelled, remove the current node monitor from the
      // cancellations set, so we can check at the end that the cancellations set is empty to
      // ensure that the cancellations returned were correct.
      if (cancellations != null) {
        assertTrue(cancellations.contains(hostPort));
        cancellations.remove(hostPort);
      }

      for (int taskIndex = 0; taskIndex < entry.getValue().numTasks;
          taskIndex++) {
        ++numGetTasks;
        List<TTaskLaunchSpec> specs = taskPlacer.assignTask(hostPort);
        assertTrue(specs != null);
        if (numGetTasks == 0) {
          // If this is the first attempt, we should definitely get a task.
          assertEquals(specs.size(), 1);
        }
        if (specs.size() == 1) {
          assertTrue("assignTask() returned a task after all tasks were launched!",
              unlaunchedTaskIds.size() > 0);
          TTaskLaunchSpec spec = specs.get(0);
          assertTrue(spec != null);
          assertTrue(unlaunchedTaskIds.contains(spec.getTaskId()));
          unlaunchedTaskIds.remove(spec.getTaskId());
        }

        // If all of the tasks have been launched, cancel the outstanding ones.
        if (unlaunchedTaskIds.size() == 0 && cancellations == null) {
          assertTrue(taskPlacer.allTasksPlaced());
          cancellations = new HashSet<THostPort>(
              taskPlacer.getOutstandingNodeMonitorsForCancellation());
          assertTrue(cancellations.size() > 0);

          // If not all of the reservations at the current node monitor have been replied to yet,
          // there should be a cancellation for the current node monitor.
          if (taskIndex < entry.getValue().numTasks - 1) {
            assertTrue(cancellations.contains(hostPort));
            cancellations.remove(hostPort);
          }

          // Another cancellation should lead to no outstanding node monitors.
          Set<THostPort> secondCancellations = new HashSet<THostPort>(
              taskPlacer.getOutstandingNodeMonitorsForCancellation());
          assertEquals(secondCancellations.size(), 0);

          // Keep going after the cancellation, since this can happen if a getTask() is in flight
          // when the original cancellation occurs.
        }
      }
    }
  }

  /**
   * Creates a scheduling request with 3 constrained tasks and 1 unconstrained task, and sanity
   * checks the result of getEnqueueTaskReservationsRequests(). Then verifies the result of
   * assignTask().
   */
  @Test
  public void testConstrainedAndUnconstrainedTasks() {
    List<InetSocketAddress> preferredNodes = new ArrayList<InetSocketAddress>();
    preferredNodes.add(new InetSocketAddress("123.4.5.6", 1));
    preferredNodes.add(new InetSocketAddress("3.4.5.6", 56));
    preferredNodes.add(new InetSocketAddress("4.2.3.67", 89));
    preferredNodes.add(new InetSocketAddress("4.1.23.45", 100));
    preferredNodes.add(new InetSocketAddress("9.0.0.0", 204));

    // Create 3 tasks with overlapping placement preferences (all tasks include nodes[2] in their
    // placement preferences).
    ByteBuffer message = ByteBuffer.allocate(1);
    List<TTaskSpec> tasks = new ArrayList<TTaskSpec>();
    final int NUM_TASKS = 3;
    Set<String> constrainedTaskIds = new HashSet<String>();
    for (int i = 0; i < NUM_TASKS; ++i) {
      TPlacementPreference placementPreference = new TPlacementPreference();
      for (int j = i; j < i + 3; ++j) {
        placementPreference.addToNodes(preferredNodes.get(j).getAddress().getHostAddress());
      }
      String id = "constrained test task " + i;
      constrainedTaskIds.add(id);
      tasks.add(new TTaskSpec(id, placementPreference, message));
    }

    String unconstrainedTaskId = "unconstrained test task";
    tasks.add(new TTaskSpec(unconstrainedTaskId, null, message));

    TUserGroupInfo user = new TUserGroupInfo(USER, GROUP, PRIORITY);
    TSchedulingRequest schedulingRequest = new TSchedulingRequest(APP_ID, tasks, user);

    // Create list of available backend nodes, with some additional nodes in additional to the
    // preferred ones.
    List<InetSocketAddress> backendNodes = new ArrayList<InetSocketAddress>(preferredNodes);
    backendNodes.add(new InetSocketAddress("1.2.3.4", 345));
    backendNodes.add(new InetSocketAddress("5.6.7.81", 2000));

    for (int i = 0; i < ITERATIONS; ++i) {
      ConstrainedTaskPlacer taskPlacer = new ConstrainedTaskPlacer(REQUEST_ID, 2);

      Map<InetSocketAddress, TEnqueueTaskReservationsRequest> requests =
          taskPlacer.getEnqueueTaskReservationsRequests(schedulingRequest, REQUEST_ID,
                                                        backendNodes, SCHEDULER_ADDRESS);

      // Sanity check the list of requests.
      int totalEnqueuedTasks = 0;
      /* Number of requests not in the list of preferred nodes (expect at most 2, corresponding
       * to the two reservations for the unconstrained task). */
      int nonPreferredNodes = 0;
      for (Entry<InetSocketAddress, TEnqueueTaskReservationsRequest> entry : requests.entrySet()) {
        if (!preferredNodes.contains(entry.getKey())) {
          nonPreferredNodes++;
          assertTrue("Expect at most two entrys to be outside the set of preferred nodes",
              nonPreferredNodes <= 2);
        }
        TEnqueueTaskReservationsRequest request = entry.getValue();
        assertEquals(request.getAppId(), APP_ID);
        assertEquals(request.getUser(), user);
        assertEquals(request.getRequestId(), REQUEST_ID);
        assertEquals(request.getSchedulerAddress(), SCHEDULER_ADDRESS);
        assertTrue(request.getNumTasks() > 0);
        totalEnqueuedTasks += request.getNumTasks();
      }
      // Expect two choices for each of the 4 tasks.
      assertEquals(totalEnqueuedTasks, 8);

      sanityCheckGetTasksAndCancellation(schedulingRequest, taskPlacer, requests);
    }

  }

}
TOP

Related Classes of edu.berkeley.sparrow.daemon.scheduler.TestConstrainedTaskPlacer

TOP
Copyright © 2018 www.massapi.com. All rights reserved.
All source code are property of their respective owners. Java is a trademark of Sun Microsystems, Inc and owned by ORACLE Inc. Contact coftware#gmail.com.