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

Example 6 with TaskModel

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

the class TestGroupByContainerCount method testGroupTaskCountDecrease.

@Test
public void testGroupTaskCountDecrease() {
    int taskCount = 3;
    Set<TaskModel> taskModels = generateTaskModels(taskCount);
    Set<ContainerModel> prevContainers = new GroupByContainerCount(3).group(generateTaskModels(taskCount + 1));
    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).deleteTaskContainerMappings(anyCollection());
}

Example 7 with TaskModel

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

the class TestGroupByContainerCount method testBalancerThrowsOnNonIntegerContainerIds.

@Test(expected = SamzaException.class)
public void testBalancerThrowsOnNonIntegerContainerIds() {
    Set<TaskModel> taskModels = generateTaskModels(3);
    Set<ContainerModel> prevContainers = new HashSet<>();
    taskModels.forEach(model -> {
        prevContainers.add(new ContainerModel(UUID.randomUUID().toString(), -1, Collections.singletonMap(model.getTaskName(), model)));
    });
    Map<String, String> prevTaskToContainerMapping = generateTaskContainerMapping(prevContainers);
    when(taskAssignmentManager.readTaskAssignment()).thenReturn(prevTaskToContainerMapping);
    //Should throw
    new GroupByContainerCount(3).balance(taskModels, localityManager);
}

Example 8 with TaskModel

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

the class TestGroupByContainerCount method testGroupHappyPath.

@Test
public void testGroupHappyPath() {
    Set<TaskModel> taskModels = generateTaskModels(5);
    Set<ContainerModel> containers = new GroupByContainerCount(2).group(taskModels);
    Map<String, ContainerModel> containersMap = new HashMap<>();
    for (ContainerModel container : containers) {
        containersMap.put(container.getProcessorId(), container);
    }
    assertEquals(2, containers.size());
    ContainerModel container0 = containersMap.get("0");
    ContainerModel container1 = containersMap.get("1");
    assertNotNull(container0);
    assertNotNull(container1);
    assertEquals("0", container0.getProcessorId());
    assertEquals("1", 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 9 with TaskModel

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

the class TestGroupByContainerCount method testBalancerOldContainerCountOne.

@Test
public void testBalancerOldContainerCountOne() {
    Set<TaskModel> taskModels = generateTaskModels(3);
    Set<ContainerModel> prevContainers = new GroupByContainerCount(1).group(taskModels);
    Map<String, String> prevTaskToContainerMapping = generateTaskContainerMapping(prevContainers);
    when(taskAssignmentManager.readTaskAssignment()).thenReturn(prevTaskToContainerMapping);
    Set<ContainerModel> containers = new GroupByContainerCount(3).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(3, containers.size());
    ContainerModel container0 = containersMap.get("0");
    ContainerModel container1 = containersMap.get("1");
    ContainerModel container2 = containersMap.get("2");
    assertNotNull(container0);
    assertNotNull(container1);
    assertNotNull(container2);
    assertEquals("0", container0.getProcessorId());
    assertEquals("1", container1.getProcessorId());
    assertEquals("2", container2.getProcessorId());
    assertEquals(1, container0.getTasks().size());
    assertEquals(1, container1.getTasks().size());
    assertEquals(1, container2.getTasks().size());
    assertTrue(container0.getTasks().containsKey(getTaskName(0)));
    assertTrue(container1.getTasks().containsKey(getTaskName(1)));
    assertTrue(container2.getTasks().containsKey(getTaskName(2)));
    // Verify task mappings are saved
    verify(taskAssignmentManager).writeTaskContainerMapping(getTaskName(0).getTaskName(), "0");
    verify(taskAssignmentManager).writeTaskContainerMapping(getTaskName(1).getTaskName(), "1");
    verify(taskAssignmentManager).writeTaskContainerMapping(getTaskName(2).getTaskName(), "2");
    verify(taskAssignmentManager, never()).deleteTaskContainerMappings(anyCollection());
}

Example 10 with TaskModel

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

the class TestGroupByContainerCount method testBalancerEmptyTaskMapping.

@Test
public void testBalancerEmptyTaskMapping() {
    Set<TaskModel> taskModels = generateTaskModels(3);
    when(taskAssignmentManager.readTaskAssignment()).thenReturn(new HashMap<String, String>());
    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());
}