Java Code Examples for org.apache.helix.task.TaskDriver#getJobContext()

@GET
@Path("{jobName}")
public Response getJob(@PathParam("clusterId") String clusterId,
    @PathParam("workflowName") String workflowName, @PathParam("jobName") String jobName) {
  TaskDriver driver = getTaskDriver(clusterId);
  Map<String, ZNRecord> jobMap = new HashMap<>();


  JobConfig jobConfig = driver.getJobConfig(jobName);
  if (jobConfig != null) {
    jobMap.put(JobProperties.JobConfig.name(), jobConfig.getRecord());
  } else {
    return badRequest(String.format("Job config for %s does not exists", jobName));
  }

  JobContext jobContext =
      driver.getJobContext(jobName);
  jobMap.put(JobProperties.JobContext.name(), null);

  if (jobContext != null) {
    jobMap.put(JobProperties.JobContext.name(), jobContext.getRecord());
  }

  return JSONRepresentation(jobMap);
}
 

private Integer getPartitionForHelixTask(TaskDriver taskDriver) {
  //Get Helix partition id for this task
  JobContext jobContext = taskDriver.getJobContext(this.helixJobId);
  if (jobContext != null) {
    return jobContext.getTaskIdPartitionMap().get(this.helixTaskId);
  }
  return null;
}
 

@GET
@Path("{jobName}/context")
public Response getJobContext(@PathParam("clusterId") String clusterId,
    @PathParam("workflowName") String workflowName, @PathParam("jobName") String jobName) {
  TaskDriver driver = getTaskDriver(clusterId);

  JobContext jobContext =
      driver.getJobContext(jobName);
  if (jobContext != null) {
    return JSONRepresentation(jobContext.getRecord());
  }
  return badRequest("Job context for " + jobName + " does not exists");
}
 

StringRepresentation getHostedEntitiesRepresentation(String clusterName, String jobQueueName,
    String jobName) throws Exception {

  ZkClient zkClient =
      ResourceUtil.getAttributeFromCtx(getContext(), ResourceUtil.ContextKey.ZKCLIENT);
  HelixDataAccessor accessor =
      ClusterRepresentationUtil.getClusterDataAccessor(zkClient, clusterName);
  PropertyKey.Builder keyBuilder = accessor.keyBuilder();

  // Get job queue config
  String namespacedJobName = TaskUtil.getNamespacedJobName(jobQueueName, jobName);
  HelixProperty jobConfig = accessor.getProperty(keyBuilder.resourceConfig(namespacedJobName));

  TaskDriver taskDriver = new TaskDriver(zkClient, clusterName);

  // Get job queue context
  JobContext ctx = taskDriver.getJobContext(namespacedJobName);

  // Create the result
  ZNRecord hostedEntitiesRecord = new ZNRecord(namespacedJobName);
  if (jobConfig != null) {
    hostedEntitiesRecord.merge(jobConfig.getRecord());
  }
  if (ctx != null) {
    hostedEntitiesRecord.merge(ctx.getRecord());
  }

  StringRepresentation representation =
      new StringRepresentation(ClusterRepresentationUtil.ZNRecordToJson(hostedEntitiesRecord),
          MediaType.APPLICATION_JSON);

  return representation;
}
 

public void testJobCleanup() throws Exception {
  final ConcurrentHashMap<String, Boolean> runningMap = new ConcurrentHashMap<>();

  final Properties properties = generateJobProperties(this.baseConfig, "3", "_1504201348473");
  final GobblinHelixJobLauncher gobblinHelixJobLauncher =
      new GobblinHelixJobLauncher(properties, this.helixManager, this.appWorkDir, ImmutableList.<Tag<?>>of(), runningMap,
          java.util.Optional.empty());

  final Properties properties2 = generateJobProperties(this.baseConfig, "33", "_1504201348474");
  final GobblinHelixJobLauncher gobblinHelixJobLauncher2 =
      new GobblinHelixJobLauncher(properties2, this.helixManager, this.appWorkDir, ImmutableList.<Tag<?>>of(), runningMap,
          java.util.Optional.empty());

  gobblinHelixJobLauncher.launchJob(null);
  gobblinHelixJobLauncher2.launchJob(null);

  final TaskDriver taskDriver = new TaskDriver(this.helixManager);

  final String jobIdKey1 = properties.getProperty(ConfigurationKeys.JOB_ID_KEY);
  final String jobIdKey2 = properties2.getProperty(ConfigurationKeys.JOB_ID_KEY);

  org.apache.helix.task.JobContext jobContext1 = taskDriver.getJobContext(jobIdKey1);
  org.apache.helix.task.JobContext jobContext2 = taskDriver.getJobContext(jobIdKey2);

  waitForWorkFlowStartup(taskDriver, jobIdKey1);
  waitForWorkFlowStartup(taskDriver, jobIdKey2);

  // job context should be present until close
  Assert.assertNotNull(jobContext1);
  Assert.assertNotNull(jobContext2);

  gobblinHelixJobLauncher.close();

  // workflow deleted asynchronously after close
  waitForWorkFlowCleanup(taskDriver, jobIdKey1);

  jobContext1 = taskDriver.getJobContext(jobIdKey1);

  // job context should have been deleted
  Assert.assertNull(jobContext1);

  // workflow should have been deleted
  WorkflowConfig workflowConfig  = taskDriver.getWorkflowConfig(jobIdKey1);
  Assert.assertNull(workflowConfig);

  WorkflowContext workflowContext = taskDriver.getWorkflowContext(jobIdKey1);
  Assert.assertNull(workflowContext);

  // second workflow with shared prefix should not be deleted when the first workflow is cleaned up
  workflowConfig  = taskDriver.getWorkflowConfig(jobIdKey2);
  Assert.assertNotNull(workflowConfig);

  gobblinHelixJobLauncher2.close();

  // workflow deleted asynchronously after close
  waitForWorkFlowCleanup(taskDriver, jobIdKey2);

  workflowConfig  = taskDriver.getWorkflowConfig(jobIdKey2);
  Assert.assertNull(workflowConfig);

  // check that workunit and taskstate directory for the job are cleaned up
  final File workunitsDir =
      new File(this.appWorkDir + File.separator + GobblinClusterConfigurationKeys.INPUT_WORK_UNIT_DIR_NAME
      + File.separator + jobIdKey1);

  final File taskstatesDir =
      new File(this.appWorkDir + File.separator + GobblinClusterConfigurationKeys.OUTPUT_TASK_STATE_DIR_NAME
          + File.separator + jobIdKey1);

  Assert.assertFalse(workunitsDir.exists());
  Assert.assertFalse(taskstatesDir.exists());

  // check that job.state file is cleaned up
  final File jobStateFile = new File(GobblinClusterUtils.getJobStateFilePath(true, this.appWorkDir, jobIdKey1).toString());

  Assert.assertFalse(jobStateFile.exists());
}
 

private static void list(TaskDriver taskDriver, String workflow) {
  WorkflowConfig wCfg = taskDriver.getWorkflowConfig(workflow);
  if (wCfg == null) {
    LOG.error("Workflow " + workflow + " does not exist!");
    return;
  }
  WorkflowContext wCtx = taskDriver.getWorkflowContext(workflow);

  LOG.info("Workflow " + workflow + " consists of the following tasks: " + wCfg.getJobDag()
      .getAllNodes());
  String workflowState =
      (wCtx != null) ? wCtx.getWorkflowState().name() : TaskState.NOT_STARTED.name();
  LOG.info("Current state of workflow is " + workflowState);
  LOG.info("Job states are: ");
  LOG.info("-------");
  for (String job : wCfg.getJobDag().getAllNodes()) {
    TaskState jobState = (wCtx != null) ? wCtx.getJobState(job) : TaskState.NOT_STARTED;
    LOG.info("Job " + job + " is " + jobState);

    // fetch job information
    JobConfig jCfg = taskDriver.getJobConfig(job);
    JobContext jCtx = taskDriver.getJobContext(job);
    if (jCfg == null || jCtx == null) {
      LOG.info("-------");
      continue;
    }

    // calculate taskPartitions
    List<Integer> partitions = Lists.newArrayList(jCtx.getPartitionSet());
    Collections.sort(partitions);

    // report status
    for (Integer partition : partitions) {
      String taskId = jCtx.getTaskIdForPartition(partition);
      taskId = (taskId != null) ? taskId : jCtx.getTargetForPartition(partition);
      LOG.info("Task: " + taskId);
      TaskConfig taskConfig = jCfg.getTaskConfig(taskId);
      if (taskConfig != null) {
        LOG.info("Configuration: " + taskConfig.getConfigMap());
      }
      TaskPartitionState state = jCtx.getPartitionState(partition);
      state = (state != null) ? state : TaskPartitionState.INIT;
      LOG.info("State: " + state);
      String assignedParticipant = jCtx.getAssignedParticipant(partition);
      if (assignedParticipant != null) {
        LOG.info("Assigned participant: " + assignedParticipant);
      }
      LOG.info("-------");
    }
    LOG.info("-------");
  }
}