Examples with WorkUnitState org.apache.gobblin.configuration.WorkUnitState used on opensource projects

Example 36 with WorkUnitState

use of org.apache.gobblin.configuration.WorkUnitState in project incubator-gobblin by apache.

the class AbstractJobLauncher method launchJob.

@Override
public void launchJob(JobListener jobListener) throws JobException {
    String jobId = this.jobContext.getJobId();
    final JobState jobState = this.jobContext.getJobState();
    try {
        MDC.put(ConfigurationKeys.JOB_NAME_KEY, this.jobContext.getJobName());
        MDC.put(ConfigurationKeys.JOB_KEY_KEY, this.jobContext.getJobKey());
        TimingEvent launchJobTimer = this.eventSubmitter.getTimingEvent(TimingEvent.LauncherTimings.FULL_JOB_EXECUTION);
        try (Closer closer = Closer.create()) {
            closer.register(this.jobContext);
            notifyListeners(this.jobContext, jobListener, TimingEvent.LauncherTimings.JOB_PREPARE, new JobListenerAction() {

                @Override
                public void apply(JobListener jobListener, JobContext jobContext) throws Exception {
                    jobListener.onJobPrepare(jobContext);
                }
            });
            if (this.jobContext.getSemantics() == DeliverySemantics.EXACTLY_ONCE) {
                // If exactly-once is used, commit sequences of the previous run must be successfully compelted
                // before this run can make progress.
                executeUnfinishedCommitSequences(jobState.getJobName());
            }
            TimingEvent workUnitsCreationTimer = this.eventSubmitter.getTimingEvent(TimingEvent.LauncherTimings.WORK_UNITS_CREATION);
            Source<?, ?> source = this.jobContext.getSource();
            WorkUnitStream workUnitStream;
            if (source instanceof WorkUnitStreamSource) {
                workUnitStream = ((WorkUnitStreamSource) source).getWorkunitStream(jobState);
            } else {
                workUnitStream = new BasicWorkUnitStream.Builder(source.getWorkunits(jobState)).build();
            }
            workUnitsCreationTimer.stop(this.eventMetadataGenerator.getMetadata(this.jobContext, EventName.WORK_UNITS_CREATION));
            // The absence means there is something wrong getting the work units
            if (workUnitStream == null || workUnitStream.getWorkUnits() == null) {
                this.eventSubmitter.submit(JobEvent.WORK_UNITS_MISSING);
                jobState.setState(JobState.RunningState.FAILED);
                throw new JobException("Failed to get work units for job " + jobId);
            }
            // No work unit to run
            if (!workUnitStream.getWorkUnits().hasNext()) {
                this.eventSubmitter.submit(JobEvent.WORK_UNITS_EMPTY);
                LOG.warn("No work units have been created for job " + jobId);
                jobState.setState(JobState.RunningState.COMMITTED);
                notifyListeners(this.jobContext, jobListener, TimingEvent.LauncherTimings.JOB_COMPLETE, new JobListenerAction() {

                    @Override
                    public void apply(JobListener jobListener, JobContext jobContext) throws Exception {
                        jobListener.onJobCompletion(jobContext);
                    }
                });
                return;
            }
            // Initialize writer and converter(s)
            closer.register(WriterInitializerFactory.newInstace(jobState, workUnitStream)).initialize();
            closer.register(ConverterInitializerFactory.newInstance(jobState, workUnitStream)).initialize();
            TimingEvent stagingDataCleanTimer = this.eventSubmitter.getTimingEvent(TimingEvent.RunJobTimings.MR_STAGING_DATA_CLEAN);
            // Cleanup left-over staging data possibly from the previous run. This is particularly
            // important if the current batch of WorkUnits include failed WorkUnits from the previous
            // run which may still have left-over staging data not cleaned up yet.
            cleanLeftoverStagingData(workUnitStream, jobState);
            stagingDataCleanTimer.stop(this.eventMetadataGenerator.getMetadata(this.jobContext, EventName.MR_STAGING_DATA_CLEAN));
            long startTime = System.currentTimeMillis();
            jobState.setStartTime(startTime);
            jobState.setState(JobState.RunningState.RUNNING);
            try {
                LOG.info("Starting job " + jobId);
                notifyListeners(this.jobContext, jobListener, TimingEvent.LauncherTimings.JOB_START, new JobListenerAction() {

                    @Override
                    public void apply(JobListener jobListener, JobContext jobContext) throws Exception {
                        jobListener.onJobStart(jobContext);
                    }
                });
                TimingEvent workUnitsPreparationTimer = this.eventSubmitter.getTimingEvent(TimingEvent.LauncherTimings.WORK_UNITS_PREPARATION);
                // Add task ids
                workUnitStream = prepareWorkUnits(workUnitStream, jobState);
                // Remove skipped workUnits from the list of work units to execute.
                workUnitStream = workUnitStream.filter(new SkippedWorkUnitsFilter(jobState));
                // Add surviving tasks to jobState
                workUnitStream = workUnitStream.transform(new MultiWorkUnitForEach() {

                    @Override
                    public void forWorkUnit(WorkUnit workUnit) {
                        jobState.incrementTaskCount();
                        jobState.addTaskState(new TaskState(new WorkUnitState(workUnit, jobState)));
                    }
                });
                workUnitsPreparationTimer.stop(this.eventMetadataGenerator.getMetadata(this.jobContext, EventName.WORK_UNITS_PREPARATION));
                // Write job execution info to the job history store before the job starts to run
                this.jobContext.storeJobExecutionInfo();
                TimingEvent jobRunTimer = this.eventSubmitter.getTimingEvent(TimingEvent.LauncherTimings.JOB_RUN);
                // Start the job and wait for it to finish
                runWorkUnitStream(workUnitStream);
                jobRunTimer.stop(this.eventMetadataGenerator.getMetadata(this.jobContext, EventName.JOB_RUN));
                this.eventSubmitter.submit(CaseFormat.UPPER_UNDERSCORE.to(CaseFormat.UPPER_CAMEL, "JOB_" + jobState.getState()));
                // Check and set final job jobPropsState upon job completion
                if (jobState.getState() == JobState.RunningState.CANCELLED) {
                    LOG.info(String.format("Job %s has been cancelled, aborting now", jobId));
                    return;
                }
                TimingEvent jobCommitTimer = this.eventSubmitter.getTimingEvent(TimingEvent.LauncherTimings.JOB_COMMIT);
                this.jobContext.finalizeJobStateBeforeCommit();
                this.jobContext.commit();
                postProcessJobState(jobState);
                jobCommitTimer.stop(this.eventMetadataGenerator.getMetadata(this.jobContext, EventName.JOB_COMMIT));
            } finally {
                long endTime = System.currentTimeMillis();
                jobState.setEndTime(endTime);
                jobState.setDuration(endTime - jobState.getStartTime());
            }
        } catch (Throwable t) {
            jobState.setState(JobState.RunningState.FAILED);
            String errMsg = "Failed to launch and run job " + jobId;
            LOG.error(errMsg + ": " + t, t);
        } finally {
            try {
                TimingEvent jobCleanupTimer = this.eventSubmitter.getTimingEvent(TimingEvent.LauncherTimings.JOB_CLEANUP);
                cleanupStagingData(jobState);
                jobCleanupTimer.stop(this.eventMetadataGenerator.getMetadata(this.jobContext, EventName.JOB_CLEANUP));
                // Write job execution info to the job history store upon job termination
                this.jobContext.storeJobExecutionInfo();
            } finally {
                launchJobTimer.stop(this.eventMetadataGenerator.getMetadata(this.jobContext, EventName.FULL_JOB_EXECUTION));
            }
        }
        for (JobState.DatasetState datasetState : this.jobContext.getDatasetStatesByUrns().values()) {
            // Set the overall job state to FAILED if the job failed to process any dataset
            if (datasetState.getState() == JobState.RunningState.FAILED) {
                jobState.setState(JobState.RunningState.FAILED);
                LOG.warn("At least one dataset state is FAILED. Setting job state to FAILED.");
                break;
            }
        }
        notifyListeners(this.jobContext, jobListener, TimingEvent.LauncherTimings.JOB_COMPLETE, new JobListenerAction() {

            @Override
            public void apply(JobListener jobListener, JobContext jobContext) throws Exception {
                jobListener.onJobCompletion(jobContext);
            }
        });
        if (jobState.getState() == JobState.RunningState.FAILED) {
            notifyListeners(this.jobContext, jobListener, TimingEvent.LauncherTimings.JOB_FAILED, new JobListenerAction() {

                @Override
                public void apply(JobListener jobListener, JobContext jobContext) throws Exception {
                    jobListener.onJobFailure(jobContext);
                }
            });
            throw new JobException(String.format("Job %s failed", jobId));
        }
    } finally {
        // Stop metrics reporting
        if (this.jobContext.getJobMetricsOptional().isPresent()) {
            JobMetrics.remove(jobState);
        }
        MDC.remove(ConfigurationKeys.JOB_NAME_KEY);
        MDC.remove(ConfigurationKeys.JOB_KEY_KEY);
    }
}

Example 37 with WorkUnitState

use of org.apache.gobblin.configuration.WorkUnitState in project incubator-gobblin by apache.

the class ConverterTest method testMultiOutputIterable.

@Test
public void testMultiOutputIterable() throws Exception {
    MyConverter converter = new MyConverter();
    BasicAckableForTesting ackable = new BasicAckableForTesting();
    RecordStreamWithMetadata<Integer, String> stream = new RecordStreamWithMetadata<>(Flowable.just(new RecordEnvelope<>(2)), GlobalMetadata.<String>builder().schema("schema").build()).mapRecords(r -> {
        r.addCallBack(ackable);
        return r;
    });
    List<StreamEntity<Integer>> outputRecords = Lists.newArrayList();
    converter.processStream(stream, new WorkUnitState()).getRecordStream().subscribe(outputRecords::add);
    Assert.assertEquals(outputRecords.size(), 2);
    // output record has not been acked
    Assert.assertEquals(ackable.acked, 0);
    outputRecords.get(0).ack();
    // only one output record acked, still need to ack another derived record
    Assert.assertEquals(ackable.acked, 0);
    outputRecords.get(1).ack();
    // all output records acked
    Assert.assertEquals(ackable.acked, 1);
}

Example 38 with WorkUnitState

use of org.apache.gobblin.configuration.WorkUnitState in project incubator-gobblin by apache.

the class ConverterTest method testMixedStream.

@Test
public void testMixedStream() throws Exception {
    MyConverter converter = new MyConverter();
    BasicAckableForTesting ackable = new BasicAckableForTesting();
    RecordStreamWithMetadata<Integer, String> stream = new RecordStreamWithMetadata<>(Flowable.just(new RecordEnvelope<>(1), new MyControlMessage<>()), GlobalMetadata.<String>builder().schema("schema").build()).mapRecords(r -> {
        r.addCallBack(ackable);
        return r;
    });
    List<StreamEntity<Integer>> outputRecords = Lists.newArrayList();
    converter.processStream(stream, new WorkUnitState()).getRecordStream().subscribe(outputRecords::add);
    Assert.assertEquals(outputRecords.size(), 2);
    Assert.assertEquals(((RecordEnvelope<Integer>) outputRecords.get(0)).getRecord(), new Integer(0));
    Assert.assertTrue(outputRecords.get(1) instanceof MyControlMessage);
}

Example 39 with WorkUnitState

use of org.apache.gobblin.configuration.WorkUnitState in project incubator-gobblin by apache.

the class ForkerTest method test.

@Test
public void test() throws Exception {
    Forker forker = new Forker();
    MyFlowable<StreamEntity<byte[]>> flowable = new MyFlowable<>();
    RecordStreamWithMetadata<byte[], String> stream = new RecordStreamWithMetadata<>(flowable, GlobalMetadata.<String>builder().schema("schema").build());
    WorkUnitState workUnitState = new WorkUnitState();
    workUnitState.setProp(ConfigurationKeys.FORK_BRANCHES_KEY, "3");
    Forker.ForkedStream<byte[], String> forkedStream = forker.forkStream(stream, new MyForkOperator(), workUnitState);
    Assert.assertEquals(forkedStream.getForkedStreams().size(), 3);
    Queue<StreamEntity<byte[]>> output0 = new LinkedList<>();
    forkedStream.getForkedStreams().get(0).getRecordStream().subscribe(output0::add);
    Queue<StreamEntity<byte[]>> output1 = new LinkedList<>();
    forkedStream.getForkedStreams().get(1).getRecordStream().subscribe(output1::add);
    Queue<StreamEntity<byte[]>> output2 = new LinkedList<>();
    forkedStream.getForkedStreams().get(2).getRecordStream().subscribe(output2::add);
    flowable._subscriber.onNext(new RecordEnvelope<>(new byte[] { 1, 1, 1 }));
    Assert.assertTrue(output0.poll() instanceof RecordEnvelope);
    Assert.assertTrue(output1.poll() instanceof RecordEnvelope);
    Assert.assertTrue(output2.poll() instanceof RecordEnvelope);
    flowable._subscriber.onNext(new RecordEnvelope<>(new byte[] { 1, 0, 0 }));
    Assert.assertTrue(output0.poll() instanceof RecordEnvelope);
    Assert.assertNull(output1.poll());
    Assert.assertNull(output2.poll());
    flowable._subscriber.onNext(new RecordEnvelope<>(new byte[] { 0, 1, 1 }));
    Assert.assertNull(output0.poll());
    Assert.assertTrue(output1.poll() instanceof RecordEnvelope);
    Assert.assertTrue(output2.poll() instanceof RecordEnvelope);
    flowable._subscriber.onNext(new BasicTestControlMessage<byte[]>("control"));
    Assert.assertTrue(output0.poll() instanceof BasicTestControlMessage);
    Assert.assertTrue(output1.poll() instanceof BasicTestControlMessage);
    Assert.assertTrue(output2.poll() instanceof BasicTestControlMessage);
    flowable._subscriber.onComplete();
}

Example 40 with WorkUnitState

use of org.apache.gobblin.configuration.WorkUnitState in project incubator-gobblin by apache.

the class BaseDataPublisher method publishData.

@Override
public void publishData(Collection<? extends WorkUnitState> states) throws IOException {
    // We need a Set to collect unique writer output paths as multiple tasks may belong to the same extract. Tasks that
    // belong to the same Extract will by default have the same output directory
    Set<Path> writerOutputPathsMoved = Sets.newHashSet();
    for (WorkUnitState workUnitState : states) {
        for (int branchId = 0; branchId < this.numBranches; branchId++) {
            publishMultiTaskData(workUnitState, branchId, writerOutputPathsMoved);
        }
    }
    this.parallelRunnerCloser.close();
    for (WorkUnitState workUnitState : states) {
        // Upon successfully committing the data to the final output directory, set states
        // of successful tasks to COMMITTED. leaving states of unsuccessful ones unchanged.
        // This makes sense to the COMMIT_ON_PARTIAL_SUCCESS policy.
        workUnitState.setWorkingState(WorkUnitState.WorkingState.COMMITTED);
    }
}