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

Example 1 with ContainerModel

use of org.apache.samza.job.model.ContainerModel 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 2 with ContainerModel

use of org.apache.samza.job.model.ContainerModel 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 3 with ContainerModel

use of org.apache.samza.job.model.ContainerModel 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 4 with ContainerModel

use of org.apache.samza.job.model.ContainerModel 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());
}

Example 5 with ContainerModel

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

the class TestHostAwareContainerAllocator method getJobModelManager.

private static JobModelManager getJobModelManager(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);
}