Java Code Examples for org.apache.helix.task.WorkflowConfig#getJobDag()

@Test
public void testBatchAddJobs() throws Exception {
  TaskDriver driver = new TaskDriver(_gZkClient, CLUSTER_NAME);
  driver.createQueue(new JobQueue.Builder(QUEUE_NAME).build());
  for (int i = 0; i < 10; i++) {
    _submitJobTasks.add(new SubmitJobTask(ZK_ADDR, i));
    _submitJobTasks.get(i).start();
  }

  WorkflowConfig workflowConfig = driver.getWorkflowConfig(QUEUE_NAME);
  while (workflowConfig.getJobDag().getAllNodes().size() < 100) {
    Thread.sleep(50);
    driver.getWorkflowConfig(QUEUE_NAME);
  }

  JobDag dag = workflowConfig.getJobDag();
  String currentJob = dag.getAllNodes().iterator().next();
  while (dag.getDirectChildren(currentJob).size() > 0) {
    String childJob = dag.getDirectChildren(currentJob).iterator().next();
    if (!getPrefix(currentJob).equals(getPrefix(childJob))
        && currentJob.charAt(currentJob.length() - 1) != '9') {
      Assert.fail();
    }
    currentJob = childJob;
  }
}
 

private void verifyJobNotInQueue(String queueName, String namedSpacedJobName) {
  WorkflowConfig workflowCfg = _driver.getWorkflowConfig(queueName);
  JobDag dag = workflowCfg.getJobDag();
  Assert.assertFalse(dag.getAllNodes().contains(namedSpacedJobName));
  Assert.assertFalse(dag.getChildrenToParents().containsKey(namedSpacedJobName));
  Assert.assertFalse(dag.getParentsToChildren().containsKey(namedSpacedJobName));
}
 

@GET
@Path("{workflowId}")
public Response getWorkflow(@PathParam("clusterId") String clusterId,
    @PathParam("workflowId") String workflowId) {
  TaskDriver taskDriver = getTaskDriver(clusterId);
  WorkflowConfig workflowConfig = taskDriver.getWorkflowConfig(workflowId);
  WorkflowContext workflowContext = taskDriver.getWorkflowContext(workflowId);

  ObjectNode root = JsonNodeFactory.instance.objectNode();
  TextNode id = JsonNodeFactory.instance.textNode(workflowId);
  root.put(Properties.id.name(), id);

  ObjectNode workflowConfigNode = JsonNodeFactory.instance.objectNode();
  ObjectNode workflowContextNode = JsonNodeFactory.instance.objectNode();

  if (workflowConfig != null) {
    getWorkflowConfigNode(workflowConfigNode, workflowConfig.getRecord());
  }

  if (workflowContext != null) {
    getWorkflowContextNode(workflowContextNode, workflowContext.getRecord());
  }

  root.put(WorkflowProperties.WorkflowConfig.name(), workflowConfigNode);
  root.put(WorkflowProperties.WorkflowContext.name(), workflowContextNode);

  JobDag jobDag = workflowConfig.getJobDag();
  ArrayNode jobs = OBJECT_MAPPER.valueToTree(jobDag.getAllNodes());
  ObjectNode parentJobs = OBJECT_MAPPER.valueToTree(jobDag.getChildrenToParents());
  root.put(WorkflowProperties.Jobs.name(), jobs);
  root.put(WorkflowProperties.ParentJobs.name(), parentJobs);
  root.put(WorkflowProperties.LastScheduledTask.name(), OBJECT_MAPPER.valueToTree(taskDriver.getLastScheduledTaskExecutionInfo(workflowId)));
  return JSONRepresentation(root);
}
 

/**
 * Iterate through the workflows configured in Helix to figure out the number of required partitions
 * and request the {@link YarnService} to scale to the desired number of containers.
 */
@VisibleForTesting
void runInternal() {
  Set<String> inUseInstances = new HashSet<>();

  int numPartitions = 0;
  for (Map.Entry<String, WorkflowConfig> workFlowEntry : taskDriver.getWorkflows().entrySet()) {
    WorkflowContext workflowContext = taskDriver.getWorkflowContext(workFlowEntry.getKey());

    // Only allocate for active workflows
    if (workflowContext == null || !workflowContext.getWorkflowState().equals(TaskState.IN_PROGRESS)) {
      continue;
    }

    log.debug("Workflow name {} config {} context {}", workFlowEntry.getKey(), workFlowEntry.getValue(),
        workflowContext);

    WorkflowConfig workflowConfig = workFlowEntry.getValue();
    JobDag jobDag = workflowConfig.getJobDag();

    Set<String> jobs = jobDag.getAllNodes();

    // sum up the number of partitions
    for (String jobName : jobs) {
      JobContext jobContext = taskDriver.getJobContext(jobName);

      if (jobContext != null) {
        log.debug("JobContext {} num partitions {}", jobContext, jobContext.getPartitionSet().size());

        inUseInstances.addAll(jobContext.getPartitionSet().stream().map(jobContext::getAssignedParticipant)
            .filter(e -> e != null).collect(Collectors.toSet()));

        numPartitions += jobContext.getPartitionSet().size();
      }
    }
  }

  // Find all participants appearing in this cluster. Note that Helix instances can contain cluster-manager
  // and potentially replanner-instance.
  Set<String> allParticipants = getParticipants(HELIX_YARN_INSTANCE_NAME_PREFIX);

  // Find all joined participants not in-use for this round of inspection.
  // If idle time is beyond tolerance, mark the instance as unused by assigning timestamp as -1.
  for (String participant : allParticipants) {
    if (!inUseInstances.contains(participant)) {
      instanceIdleSince.putIfAbsent(participant, System.currentTimeMillis());
      if (!isInstanceUnused(participant)) {
        inUseInstances.add(participant);
      }
    } else {
      // A previously idle instance is now detected to be in use.
      // Remove this instance if existed in the tracking map.
      instanceIdleSince.remove(participant);
    }
  }

  // compute the target containers as a ceiling of number of partitions divided by the number of containers
  // per partition. Scale the result by a constant overprovision factor.
  int numTargetContainers = (int) Math.ceil(((double)numPartitions / this.partitionsPerContainer) * this.overProvisionFactor);

  // adjust the number of target containers based on the configured min and max container values.
  numTargetContainers = Math.max(this.minContainers, Math.min(this.maxContainers, numTargetContainers));

  slidingWindowReservoir.add(numTargetContainers);

  log.info("There are {} containers being requested", numTargetContainers);

  this.yarnService.requestTargetNumberOfContainers(slidingWindowReservoir.getMax(), inUseInstances);
}
 

@Test
public void testNamedQueue() throws Exception {
  String queueName = TestHelper.getTestMethodName();

  // Create a queue
  JobQueue queue = new JobQueue.Builder(queueName).build();
  _driver.createQueue(queue);

  // Enqueue jobs
  Set<String> master = Sets.newHashSet("MASTER");
  Set<String> slave = Sets.newHashSet("SLAVE");
  JobConfig.Builder job1 = new JobConfig.Builder().setCommand(MockTask.TASK_COMMAND)
      .setTargetResource(WorkflowGenerator.DEFAULT_TGT_DB).setTargetPartitionStates(master);
  JobConfig.Builder job2 = new JobConfig.Builder().setCommand(MockTask.TASK_COMMAND)
      .setTargetResource(WorkflowGenerator.DEFAULT_TGT_DB).setTargetPartitionStates(slave);
  _driver.enqueueJob(queueName, "masterJob", job1);
  _driver.enqueueJob(queueName, "slaveJob", job2);

  // Ensure successful completion
  String namespacedJob1 = queueName + "_masterJob";
  String namespacedJob2 = queueName + "_slaveJob";
  _driver.pollForJobState(queueName, namespacedJob1, TaskState.COMPLETED);
  _driver.pollForJobState(queueName, namespacedJob2, TaskState.COMPLETED);
  JobContext masterJobContext = _driver.getJobContext(namespacedJob1);
  JobContext slaveJobContext = _driver.getJobContext(namespacedJob2);

  // Ensure correct ordering
  long job1Finish = masterJobContext.getFinishTime();
  long job2Start = slaveJobContext.getStartTime();
  Assert.assertTrue(job2Start >= job1Finish);

  // Flush queue and check cleanup
  _driver.flushQueue(queueName);
  HelixDataAccessor accessor = _manager.getHelixDataAccessor();
  PropertyKey.Builder keyBuilder = accessor.keyBuilder();
  Assert.assertNull(accessor.getProperty(keyBuilder.idealStates(namespacedJob1)));
  Assert.assertNull(accessor.getProperty(keyBuilder.resourceConfig(namespacedJob1)));
  Assert.assertNull(accessor.getProperty(keyBuilder.idealStates(namespacedJob2)));
  Assert.assertNull(accessor.getProperty(keyBuilder.resourceConfig(namespacedJob2)));
  WorkflowConfig workflowCfg = _driver.getWorkflowConfig(queueName);
  JobDag dag = workflowCfg.getJobDag();
  Assert.assertFalse(dag.getAllNodes().contains(namespacedJob1));
  Assert.assertFalse(dag.getAllNodes().contains(namespacedJob2));
  Assert.assertFalse(dag.getChildrenToParents().containsKey(namespacedJob1));
  Assert.assertFalse(dag.getChildrenToParents().containsKey(namespacedJob2));
  Assert.assertFalse(dag.getParentsToChildren().containsKey(namespacedJob1));
  Assert.assertFalse(dag.getParentsToChildren().containsKey(namespacedJob2));
}