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

Example 1 with TaskModel

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

the class SingleContainerGrouper method group.

@Override
public Set<ContainerModel> group(Set<TaskModel> taskModels) {
    Map<TaskName, TaskModel> taskNameTaskModelMap = new HashMap<>();
    for (TaskModel taskModel : taskModels) {
        taskNameTaskModelMap.put(taskModel.getTaskName(), taskModel);
    }
    ContainerModel containerModel = new ContainerModel(containerId, -1, taskNameTaskModelMap);
    return Collections.singleton(containerModel);
}

Example 2 with TaskModel

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

the class TestSamzaObjectMapper method setup.

@Before
public void setup() throws IOException {
    Map<String, String> configMap = new HashMap<String, String>();
    Set<SystemStreamPartition> ssp = new HashSet<>();
    configMap.put("a", "b");
    Config config = new MapConfig(configMap);
    TaskName taskName = new TaskName("test");
    ssp.add(new SystemStreamPartition("foo", "bar", new Partition(1)));
    TaskModel taskModel = new TaskModel(taskName, ssp, new Partition(2));
    Map<TaskName, TaskModel> tasks = new HashMap<TaskName, TaskModel>();
    tasks.put(taskName, taskModel);
    ContainerModel containerModel = new ContainerModel("1", 1, tasks);
    Map<String, ContainerModel> containerMap = new HashMap<String, ContainerModel>();
    containerMap.put("1", containerModel);
    jobModel = new JobModel(config, containerMap);
}

Example 3 with TaskModel

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

the class TestGroupByContainerCount method testBalancerAfterContainerDecrease.

/**
   * Before:
   *  C0  C1  C2  C3
   * ----------------
   *  T0  T1  T2  T3
   *  T4  T5  T6  T7
   *  T8
   *
   * After:
   *  C0  C1
   * --------
   *  T0  T1
   *  T4  T5
   *  T8  T7
   *  T6  T3
   *  T2
   *
   *  NOTE for host affinity, it would help to have some additional logic to reassign tasks
   *  from C2 and C3 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 testBalancerAfterContainerDecrease() {
    Set<TaskModel> taskModels = generateTaskModels(9);
    Set<ContainerModel> prevContainers = new GroupByContainerCount(4).group(taskModels);
    Map<String, String> prevTaskToContainerMapping = generateTaskContainerMapping(prevContainers);
    when(taskAssignmentManager.readTaskAssignment()).thenReturn(prevTaskToContainerMapping);
    Set<ContainerModel> containers = new GroupByContainerCount(2).balance(taskModels, localityManager);
    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(5, container0.getTasks().size());
    assertEquals(4, container1.getTasks().size());
    // Tasks 0,4,8 and 1,5 should stay on the same original containers
    assertTrue(container0.getTasks().containsKey(getTaskName(0)));
    assertTrue(container0.getTasks().containsKey(getTaskName(4)));
    assertTrue(container0.getTasks().containsKey(getTaskName(8)));
    assertTrue(container1.getTasks().containsKey(getTaskName(1)));
    assertTrue(container1.getTasks().containsKey(getTaskName(5)));
    // Tasks 2,6 and 3,7 should be reassigned to the new containers.
    // Consistency is the goal with these reassignments
    assertTrue(container0.getTasks().containsKey(getTaskName(6)));
    assertTrue(container0.getTasks().containsKey(getTaskName(2)));
    assertTrue(container1.getTasks().containsKey(getTaskName(7)));
    assertTrue(container1.getTasks().containsKey(getTaskName(3)));
    // Verify task mappings are saved
    verify(taskAssignmentManager).writeTaskContainerMapping(getTaskName(0).getTaskName(), "0");
    verify(taskAssignmentManager).writeTaskContainerMapping(getTaskName(4).getTaskName(), "0");
    verify(taskAssignmentManager).writeTaskContainerMapping(getTaskName(8).getTaskName(), "0");
    verify(taskAssignmentManager).writeTaskContainerMapping(getTaskName(6).getTaskName(), "0");
    verify(taskAssignmentManager).writeTaskContainerMapping(getTaskName(2).getTaskName(), "0");
    verify(taskAssignmentManager).writeTaskContainerMapping(getTaskName(1).getTaskName(), "1");
    verify(taskAssignmentManager).writeTaskContainerMapping(getTaskName(5).getTaskName(), "1");
    verify(taskAssignmentManager).writeTaskContainerMapping(getTaskName(7).getTaskName(), "1");
    verify(taskAssignmentManager).writeTaskContainerMapping(getTaskName(3).getTaskName(), "1");
    verify(taskAssignmentManager, never()).deleteTaskContainerMappings(anyCollection());
}

Example 4 with TaskModel

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

the class TestGroupByContainerIds method testGroupHappyPath.

@Test
public void testGroupHappyPath() {
    Set<TaskModel> taskModels = generateTaskModels(5);
    Set<ContainerModel> containers = buildSimpleGrouper(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 5 with TaskModel

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

the class TestGroupByContainerIds method testGroupManyTasks.

@Test
public void testGroupManyTasks() {
    Set<TaskModel> taskModels = generateTaskModels(21);
    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(11, container0.getTasks().size());
    assertEquals(10, container1.getTasks().size());
    // NOTE: tasks are sorted lexicographically, so the container assignment
    // can seem odd, but the consistency is the key focus
    assertTrue(container0.getTasks().containsKey(getTaskName(0)));
    assertTrue(container0.getTasks().containsKey(getTaskName(10)));
    assertTrue(container0.getTasks().containsKey(getTaskName(12)));
    assertTrue(container0.getTasks().containsKey(getTaskName(14)));
    assertTrue(container0.getTasks().containsKey(getTaskName(16)));
    assertTrue(container0.getTasks().containsKey(getTaskName(18)));
    assertTrue(container0.getTasks().containsKey(getTaskName(2)));
    assertTrue(container0.getTasks().containsKey(getTaskName(3)));
    assertTrue(container0.getTasks().containsKey(getTaskName(5)));
    assertTrue(container0.getTasks().containsKey(getTaskName(7)));
    assertTrue(container0.getTasks().containsKey(getTaskName(9)));
    assertTrue(container1.getTasks().containsKey(getTaskName(1)));
    assertTrue(container1.getTasks().containsKey(getTaskName(11)));
    assertTrue(container1.getTasks().containsKey(getTaskName(13)));
    assertTrue(container1.getTasks().containsKey(getTaskName(15)));
    assertTrue(container1.getTasks().containsKey(getTaskName(17)));
    assertTrue(container1.getTasks().containsKey(getTaskName(19)));
    assertTrue(container1.getTasks().containsKey(getTaskName(20)));
    assertTrue(container1.getTasks().containsKey(getTaskName(4)));
    assertTrue(container1.getTasks().containsKey(getTaskName(6)));
    assertTrue(container1.getTasks().containsKey(getTaskName(8)));
}