Examples with TaskModel org.apache.samza.job.model.TaskModel used on opensource projects

Example 16 with TaskModel

use of org.apache.samza.job.model.TaskModel in project samza by apache.

the class TestContainerAllocator method getJobModelReader.

private static JobModelManager getJobModelReader(int containerCount) {
    //Ideally, the JobModelReader should be constructed independent of HttpServer.
    //That way it becomes easier to mock objects. Save it for later.
    HttpServer server = new MockHttpServer("/", 7777, null, new ServletHolder(DefaultServlet.class));
    Map<String, ContainerModel> containers = new java.util.HashMap<>();
    for (int i = 0; i < containerCount; i++) {
        ContainerModel container = new ContainerModel(String.valueOf(i), i, new HashMap<TaskName, TaskModel>());
        containers.put(String.valueOf(i), container);
    }
    JobModel jobModel = new JobModel(getConfig(), containers);
    return new JobModelManager(jobModel, server, null);
}

Example 17 with TaskModel

use of org.apache.samza.job.model.TaskModel in project samza by apache.

the class TestContainerProcessManager method getCoordinator.

private JobModelManager getCoordinator(int containerCount) {
    Map<String, ContainerModel> containers = new java.util.HashMap<>();
    for (int i = 0; i < containerCount; i++) {
        ContainerModel container = new ContainerModel(String.valueOf(i), i, new HashMap<TaskName, TaskModel>());
        containers.put(String.valueOf(i), container);
    }
    Map<String, Map<String, String>> localityMap = new HashMap<>();
    localityMap.put("0", new HashMap<String, String>() {

        {
            put(SetContainerHostMapping.HOST_KEY, "abc");
        }
    });
    LocalityManager mockLocalityManager = mock(LocalityManager.class);
    when(mockLocalityManager.readContainerLocality()).thenReturn(localityMap);
    JobModel jobModel = new JobModel(getConfig(), containers, mockLocalityManager);
    JobModelManager.jobModelRef().getAndSet(jobModel);
    return new JobModelManager(jobModel, this.server, null);
}

Example 18 with TaskModel

use of org.apache.samza.job.model.TaskModel in project samza by apache.

the class TestGroupByContainerIds method testGroupHappyPathWithListOfContainers.

@Test
public void testGroupHappyPathWithListOfContainers() {
    Set<TaskModel> taskModels = generateTaskModels(5);
    List<String> containerIds = new ArrayList<String>() {

        {
            add("4");
            add("2");
        }
    };
    Set<ContainerModel> containers = buildSimpleGrouper().group(taskModels, containerIds);
    Map<String, ContainerModel> containersMap = new HashMap<>();
    for (ContainerModel container : containers) {
        containersMap.put(container.getProcessorId(), container);
    }
    assertEquals(2, containers.size());
    ContainerModel container0 = containersMap.get("4");
    ContainerModel container1 = containersMap.get("2");
    assertNotNull(container0);
    assertNotNull(container1);
    assertEquals("4", container0.getProcessorId());
    assertEquals("2", container1.getProcessorId());
    assertEquals(3, container0.getTasks().size());
    assertEquals(2, container1.getTasks().size());
    assertTrue(container0.getTasks().containsKey(getTaskName(0)));
    assertTrue(container0.getTasks().containsKey(getTaskName(2)));
    assertTrue(container0.getTasks().containsKey(getTaskName(4)));
    assertTrue(container1.getTasks().containsKey(getTaskName(1)));
    assertTrue(container1.getTasks().containsKey(getTaskName(3)));
}

Example 19 with TaskModel

use of org.apache.samza.job.model.TaskModel in project samza by apache.

the class TestGroupByContainerCount method testBalancerNewContainerCountOne.

@Test
public void testBalancerNewContainerCountOne() {
    Set<TaskModel> taskModels = generateTaskModels(3);
    Set<ContainerModel> prevContainers = new GroupByContainerCount(3).group(taskModels);
    Map<String, String> prevTaskToContainerMapping = generateTaskContainerMapping(prevContainers);
    when(taskAssignmentManager.readTaskAssignment()).thenReturn(prevTaskToContainerMapping);
    Set<ContainerModel> containers = new GroupByContainerCount(1).balance(taskModels, localityManager);
    // Results should be the same as calling group
    Map<String, ContainerModel> containersMap = new HashMap<>();
    for (ContainerModel container : containers) {
        containersMap.put(container.getProcessorId(), container);
    }
    assertEquals(1, containers.size());
    ContainerModel container0 = containersMap.get("0");
    assertNotNull(container0);
    assertEquals("0", container0.getProcessorId());
    assertEquals(3, container0.getTasks().size());
    assertTrue(container0.getTasks().containsKey(getTaskName(0)));
    assertTrue(container0.getTasks().containsKey(getTaskName(1)));
    assertTrue(container0.getTasks().containsKey(getTaskName(2)));
    verify(taskAssignmentManager).writeTaskContainerMapping(getTaskName(0).getTaskName(), "0");
    verify(taskAssignmentManager).writeTaskContainerMapping(getTaskName(1).getTaskName(), "0");
    verify(taskAssignmentManager).writeTaskContainerMapping(getTaskName(2).getTaskName(), "0");
    verify(taskAssignmentManager, never()).deleteTaskContainerMappings(anyCollection());
}

Example 20 with TaskModel

use of org.apache.samza.job.model.TaskModel in project samza by apache.

the class TestGroupByContainerCount method testBalancerAfterContainerIncrease.

/**
   * Before:
   *  C0  C1
   * --------
   *  T0  T1
   *  T2  T3
   *  T4  T5
   *  T6  T7
   *  T8
   *
   * After:
   *  C0  C1  C2  C3
   * ----------------
   *  T0  T1  T6  T5
   *  T2  T3  T8  T7
   *  T4
   *
   *  NOTE for host affinity, it would help to have some additional logic to reassign tasks
   *  from C0 and C1 to containers that were on the same respective hosts, it wasn't implemented
   *  because the scenario is infrequent, the benefits are not guaranteed, and the code complexity
   *  wasn't worth it. It certainly could be implemented in the future.
   */
@Test
public void testBalancerAfterContainerIncrease() {
    Set<TaskModel> taskModels = generateTaskModels(9);
    Set<ContainerModel> prevContainers = new GroupByContainerCount(2).group(taskModels);
    Map<String, String> prevTaskToContainerMapping = generateTaskContainerMapping(prevContainers);
    when(taskAssignmentManager.readTaskAssignment()).thenReturn(prevTaskToContainerMapping);
    Set<ContainerModel> containers = new GroupByContainerCount(4).balance(taskModels, localityManager);
    Map<String, ContainerModel> containersMap = new HashMap<>();
    for (ContainerModel container : containers) {
        containersMap.put(container.getProcessorId(), container);
    }
    assertEquals(4, containers.size());
    ContainerModel container0 = containersMap.get("0");
    ContainerModel container1 = containersMap.get("1");
    ContainerModel container2 = containersMap.get("2");
    ContainerModel container3 = containersMap.get("3");
    assertNotNull(container0);
    assertNotNull(container1);
    assertNotNull(container2);
    assertNotNull(container3);
    assertEquals("0", container0.getProcessorId());
    assertEquals("1", container1.getProcessorId());
    assertEquals(3, container0.getTasks().size());
    assertEquals(2, container1.getTasks().size());
    assertEquals(2, container2.getTasks().size());
    assertEquals(2, container3.getTasks().size());
    // Tasks 0-4 should stay on the same original containers
    assertTrue(container0.getTasks().containsKey(getTaskName(0)));
    assertTrue(container0.getTasks().containsKey(getTaskName(2)));
    assertTrue(container0.getTasks().containsKey(getTaskName(4)));
    assertTrue(container1.getTasks().containsKey(getTaskName(1)));
    assertTrue(container1.getTasks().containsKey(getTaskName(3)));
    // Tasks 5-8 should be reassigned to the new containers.
    // Consistency is the goal with these reassignments
    assertTrue(container2.getTasks().containsKey(getTaskName(8)));
    assertTrue(container2.getTasks().containsKey(getTaskName(6)));
    assertTrue(container3.getTasks().containsKey(getTaskName(5)));
    assertTrue(container3.getTasks().containsKey(getTaskName(7)));
    // Verify task mappings are saved
    verify(taskAssignmentManager).writeTaskContainerMapping(getTaskName(0).getTaskName(), "0");
    verify(taskAssignmentManager).writeTaskContainerMapping(getTaskName(2).getTaskName(), "0");
    verify(taskAssignmentManager).writeTaskContainerMapping(getTaskName(4).getTaskName(), "0");
    verify(taskAssignmentManager).writeTaskContainerMapping(getTaskName(1).getTaskName(), "1");
    verify(taskAssignmentManager).writeTaskContainerMapping(getTaskName(3).getTaskName(), "1");
    verify(taskAssignmentManager).writeTaskContainerMapping(getTaskName(8).getTaskName(), "2");
    verify(taskAssignmentManager).writeTaskContainerMapping(getTaskName(6).getTaskName(), "2");
    verify(taskAssignmentManager).writeTaskContainerMapping(getTaskName(5).getTaskName(), "3");
    verify(taskAssignmentManager).writeTaskContainerMapping(getTaskName(7).getTaskName(), "3");
    verify(taskAssignmentManager, never()).deleteTaskContainerMappings(anyCollection());
}