arunpandianp commented on code in PR #38919:
URL: https://github.com/apache/beam/pull/38919#discussion_r3641328261
##########
runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/streaming/Work.java:
##########
@@ -261,20 +337,30 @@ public void queueCommit(WorkItemCommitRequest
commitRequest, ComputationState co
processingContext.workCommitter().accept(Commit.create(commitRequest,
computationState, this));
}
- public WindmillStateReader createWindmillStateReader() {
- return WindmillStateReader.forWork(this);
+ public Consumer<Commit> workCommitter() {
+ return processingContext.workCommitter();
+ }
+
+ public WindmillStateReader createWindmillStateReader(Supplier<Boolean>
workIsFailed) {
+ return WindmillStateReader.forWork(this, workIsFailed);
}
@Override
public WorkId id() {
return id;
}
- public void recordGetWorkStreamLatencies(
- ImmutableList<LatencyAttribution> getWorkStreamLatencies) {
- for (LatencyAttribution latency : getWorkStreamLatencies) {
- totalDurationPerState.put(
- latency.getState(),
Duration.millis(latency.getTotalDurationMillis()));
+ public ImmutableList<LatencyAttribution> getWorkStreamLatencies() {
Review Comment:
removed getWorkStreamLatencies().
recordGetWorkStreamLatencies was initially inside the constructor, moved it
out of the constructor to reduce load on the Getworkstream thread.
recordGetWorkStreamLatencies runs in the processing thread.
##########
runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/streaming/Work.java:
##########
@@ -101,12 +111,17 @@ private Work(
// keyUniverse inside EnumMap every time.
this.totalDurationPerState = new EnumMap<>(EMPTY_ENUM_MAP);
this.id = WorkId.of(workItem);
+ this.keyGroup =
+ workItem.hasKeyGroup()
+ ? KeyGroup.create(workItem.getKeyGroup().getHigh(),
workItem.getKeyGroup().getLow())
+ : KeyGroup.DEFAULT;
this.latencyTrackingId =
Long.toHexString(workItem.getShardingKey())
+ '-'
+ Long.toHexString(workItem.getWorkToken());
this.currentState = TimedState.initialState(startTime);
this.isFailed = false;
+ this.getWorkStreamLatencies = getWorkStreamLatencies;
}
public static Work create(
Review Comment:
done.
##########
runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/windmill/work/processing/StreamingWorkScheduler.java:
##########
@@ -384,90 +361,210 @@ private ExecuteWorkResult executeWork(
stageInfo, computationState,
work.getLatencyTrackingId()));
try {
- WindmillStateReader stateReader = work.createWindmillStateReader();
- SideInputStateFetcher localSideInputStateFetcher =
-
sideInputStateFetcherFactory.createSideInputStateFetcher(work::fetchSideInput);
-
- // If the read output KVs, then we can decode Windmill's byte key into
userland
- // key object and provide it to the execution context for use with
per-key state.
- // Otherwise, we pass null.
- //
- // The coder type that will be present is:
- // WindowedValueCoder(TimerOrElementCoder(KvCoder))
- Optional<Coder<?>> keyCoder = computationWorkExecutor.keyCoder();
- @SuppressWarnings("deprecation")
- @Nullable
- final Object executionKey =
- !keyCoder.isPresent() ? null : keyCoder.get().decode(key.newInput(),
Coder.Context.OUTER);
-
- if (workItem.hasHotKeyInfo()) {
- Windmill.HotKeyInfo hotKeyInfo = workItem.getHotKeyInfo();
- Duration hotKeyAge = Duration.millis(hotKeyInfo.getHotKeyAgeUsec() /
1000);
-
- String stepName =
getShuffleTaskStepName(computationState.getMapTask());
- if (executionKey != null
- && (options.isHotKeyLoggingEnabled()
- || hasExperiment(options, "enable_hot_key_logging"))
- && keyCoder.isPresent()) {
- hotKeyLogger.logHotKeyDetection(stepName, hotKeyAge, executionKey);
- } else {
- hotKeyLogger.logHotKeyDetection(stepName, hotKeyAge);
- }
- }
+ StreamingModeExecutionContext context =
computationWorkExecutor.context();
// Blocks while executing work.
- computationWorkExecutor.executeWork(
- executionKey, work, stateReader, localSideInputStateFetcher,
outputBuilder);
+ computationWorkExecutor.executeWork(work, workExecutor, handle,
keyTransitionListener);
+
+ List<Work> workBatch;
+ List<Windmill.WorkItemCommitRequest> workItemCommits;
+ Map<Long, Pair<Instant, Runnable>> finalizationCallbacks;
+ long stateBytesRead;
+ {
+ if (context.workIsFailed()) {
+ throw new
WorkItemCancelledException(work.getWorkItem().getShardingKey());
+ }
+ context.flushState();
- if (work.isFailed()) {
- throw new WorkItemCancelledException(workItem.getShardingKey());
- }
+ // Retrieve executed works, work item commits, and accumulated
callbacks from execution
+ // context
+ workBatch = context.getExecutedWorks();
+ workItemCommits = context.getWorkItemCommits();
+ finalizationCallbacks = context.getFinalizationCallbacks();
+ stateBytesRead = context.getStateBytesRead();
- // Reports source bytes processed to WorkItemCommitRequest if available.
- try {
- long sourceBytesProcessed =
- computationWorkExecutor.computeSourceBytesProcessed(
- computationState.sourceBytesProcessCounterName());
- outputBuilder.setSourceBytesProcessed(sourceBytesProcessed);
- } catch (Exception e) {
- LOG.error("{}", e.toString());
+ context.reset(); // Don't use context after this.
}
-
-
commitFinalizer.cacheCommitFinalizers(computationWorkExecutor.context().flushState());
-
// Release the execution state for another thread to use.
computationState.releaseComputationWorkExecutor(computationWorkExecutor);
computationWorkExecutor = null;
- work.setState(Work.State.COMMIT_QUEUED);
-
outputBuilder.addAllPerWorkItemLatencyAttributions(work.getLatencyAttributions(sampler));
-
return ExecuteWorkResult.create(
- outputBuilder, stateReader.getBytesRead() +
localSideInputStateFetcher.getBytesRead());
+ workBatch, workItemCommits, finalizationCallbacks, stateBytesRead);
} catch (Throwable t) {
if (computationWorkExecutor != null) {
// If processing failed due to a thrown exception, close the
executionState. Do not
// return/release the executionState back to computationState as that
will lead to this
// executionState instance being reused.
- LOG.debug("Invalidating executor after work item {} failed",
workItem.getWorkToken(), t);
+ LOG.debug(
+ "Invalidating executor after work item {} failed",
+ work.getWorkItem().getWorkToken(),
+ t);
computationWorkExecutor.invalidate();
}
-
// Re-throw the exception, it will be caught and handled by
workFailureProcessor downstream.
throw t;
}
}
+ private void handleOnlyFinalize(
+ ComputationState computationState, Work work, Windmill.WorkItem
workItem) {
+ Windmill.WorkItemCommitRequest.Builder outputBuilder =
+ initializeOutputBuilder(workItem.getKey(), workItem);
+
outputBuilder.setSourceStateUpdates(Windmill.SourceState.newBuilder().setOnlyFinalize(true));
+ work.setState(Work.State.COMMIT_QUEUED);
+ work.queueCommit(outputBuilder.build(), computationState);
+ }
+
+ private StageInfo getStageInfo(ComputationState computationState) {
+ MapTask mapTask = computationState.getMapTask();
+ return stageInfoMap.computeIfAbsent(
+ mapTask.getStageName(), s -> StageInfo.create(s,
mapTask.getSystemName()));
+ }
+
+ private void commitWorkBatch(
+ ComputationState computationState,
+ List<Work> workBatch,
+ List<Windmill.WorkItemCommitRequest> workItemCommits) {
+ if (workBatch.isEmpty()) {
+ return;
+ }
+ if (workBatch.size() > 1 || multiKeyExperimentEnabled) {
+ commitMultiKeyWorkBatch(computationState, workBatch, workItemCommits);
+ } else {
+ commitSingleKeyWork(computationState, workBatch.get(0),
workItemCommits.get(0));
+ }
+ }
+
+ private void commitMultiKeyWorkBatch(
+ ComputationState computationState,
+ List<Work> workBatch,
+ List<Windmill.WorkItemCommitRequest> workItemCommits) {
+ Windmill.MultiKeyWorkItemCommitRequest.Builder multiKeyBuilder =
+ Windmill.MultiKeyWorkItemCommitRequest.newBuilder();
+
+ Work primaryWork = workBatch.get(0);
Review Comment:
done.
##########
runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/StreamingModeExecutionContext.java:
##########
@@ -525,39 +717,196 @@ public Map<Long, Pair<Instant, Runnable>> flushState() {
throw new RuntimeException("Exception while encoding checkpoint", e);
}
sourceStateBuilder.setState(stream.toByteString());
- if
(activeReader.getCurrentSource().offsetBasedDeduplicationSupported()) {
+ if (reader.getCurrentSource().offsetBasedDeduplicationSupported()) {
byte[] offsetLimit = checkpointMark.getOffsetLimit();
if (offsetLimit.length == 0) {
throw new RuntimeException("Checkpoint offset limit must be
non-empty.");
}
sourceStateBuilder.setOffsetLimit(ByteString.copyFrom(offsetLimit));
}
}
-
outputBuilder.setSourceWatermark(WindmillTimeUtils.harnessToWindmillTimestamp(watermark));
+
builder.setSourceWatermark(WindmillTimeUtils.harnessToWindmillTimestamp(watermark));
- backlogBytes = activeReader.getSplitBacklogBytes();
+ backlogBytes = reader.getSplitBacklogBytes();
+ ByteString serializedKey = checkStateNotNull(getSerializedKey());
if (backlogBytes == UnboundedReader.BACKLOG_UNKNOWN
- &&
WorkerCustomSources.isFirstUnboundedSourceSplit(getSerializedKey())) {
+ && WorkerCustomSources.isFirstUnboundedSourceSplit(serializedKey)) {
// Only call getTotalBacklogBytes() on the first split.
- backlogBytes = activeReader.getTotalBacklogBytes();
+ backlogBytes = reader.getTotalBacklogBytes();
}
- outputBuilder.setSourceBacklogBytes(backlogBytes);
+ builder.setSourceBacklogBytes(backlogBytes);
readerCache.cacheReader(
- getComputationKey(),
- getWorkItem().getCacheToken(),
- getWorkItem().getWorkToken(),
- activeReader);
+ getComputationKey(), getWorkItem().getCacheToken(),
getWorkItem().getWorkToken(), reader);
activeReader = null;
} else if (backlogBytes != UnboundedReader.BACKLOG_UNKNOWN && backlogBytes
!= 1L) {
// If activeReader is null, we might still have backlogBytes from an
SDF. We ignore a reported
// backlogBytes of 1 since older versions of the Java SDK use this value
as a default when
// RestrictionTracker.getProgress() or GetSize() are not defined.
- outputBuilder.setSourceBacklogBytes(backlogBytes);
+ getOutputBuilder().setSourceBacklogBytes(backlogBytes);
}
- return callbacks;
+
+ this.finalizationCallbacks.putAll(callbacks);
+
+ getOutputBuilder()
+
.setSourceBytesProcessed(computeSourceBytesProcessed(sourceBytesProcessCounterName));
+ }
+
+ private final long computeSourceBytesProcessed(String
sourceBytesCounterName) {
+ if (!(workExecutor instanceof DataflowMapTaskExecutor)) {
+ return 0L;
+ }
+ HashMap<String, ElementCounter> counters =
+ ((DataflowMapTaskExecutor) workExecutor)
+ .getReadOperation()
+ .receivers[0]
+ .getOutputCounters();
+
+ return Optional.ofNullable(counters.get(sourceBytesCounterName))
+ .map(counter -> ((OutputObjectAndByteCounter)
counter).getByteCount().getAndReset())
+ .orElse(0L);
+ }
+
+ public boolean advance() throws CoderException {
+ if (!multiKeyBundleEnabled) {
+ return false;
+ }
+ if (workIsFailed()) {
+ throw new
WorkItemCancelledException(checkStateNotNull(work).getWorkItem().getShardingKey());
Review Comment:
done.
##########
runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/util/BoundedQueueExecutor.java:
##########
@@ -255,24 +262,25 @@ final class BoundedQueueExecutorWorkHandleImpl
implements BoundedQueueExecutorWorkHandle, AutoCloseable {
@GuardedBy("this")
- private int elements;
+ private final List<Work> workBatch;
@GuardedBy("this")
private long bytes;
@GuardedBy("this")
private boolean closed = false;
- private BoundedQueueExecutorWorkHandleImpl(int elements, long bytes) {
- checkArgument(elements >= 0 && bytes >= 0);
- this.elements = elements;
+ private BoundedQueueExecutorWorkHandleImpl(Work work, long bytes) {
+ checkArgument(bytes >= 0);
+ this.workBatch = new ArrayList<>();
Review Comment:
changing it to `ImmutableList.Builder` is adding more copies in merge().
ArrayList seems to be better.
##########
runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/windmill/client/commits/CompleteCommit.java:
##########
@@ -49,4 +54,6 @@ public static CompleteCommit create(
public abstract WorkId workId();
public abstract CommitStatus status();
+
+ public abstract boolean retryableFailure();
Review Comment:
Added a comment. retryableFailure is used to retry work items locally when a
subset of work fails after queuing the commit.
We could add a different value in CommitStatus to represent it, but
CommitStatus is part of the backend windmill.proto and don't want to change it.
Migrating to a different enum and mapping it to windmill.proto enum seems more
complex and unnecessary for this use case. So, I'm leaning towards keeping the
boolean.
##########
runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/streaming/ActiveWorkState.java:
##########
@@ -88,6 +88,11 @@ static ActiveWorkState
create(WindmillStateCache.ForComputation computationState
return new ActiveWorkState(new HashMap<>(), computationStateCache);
}
+ synchronized Optional<ExecutableWork> getActiveWork(ShardedKey shardedKey,
WorkId workId) {
Review Comment:
done.
##########
runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/streaming/ComputationState.java:
##########
@@ -131,6 +131,10 @@ public void
completeWorkAndScheduleNextWorkForKey(ShardedKey shardedKey, WorkId
.ifPresent(this::forceExecute);
}
+ public void reExecuteActiveWork(ShardedKey shardedKey, WorkId workId) {
Review Comment:
done.
##########
runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/WorkCancelingException.java:
##########
@@ -17,23 +17,31 @@
*/
package org.apache.beam.runners.dataflow.worker;
-import static org.junit.Assert.assertFalse;
-import static org.junit.Assert.assertTrue;
+import org.checkerframework.checker.nullness.qual.Nullable;
-import org.junit.Test;
-import org.junit.runner.RunWith;
-import org.junit.runners.JUnit4;
+/**
+ * Indicates that the work is no longer valid and should be canceled. It is
thrown as a signal for
+ * upper layers to mark the work as failed.
+ */
+public class WorkCancelingException extends RuntimeException {
Review Comment:
Renamed and added differences. once is thrown before marking work as failed
and the other is thrown after noticing that a work is failed.
##########
runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/windmill/work/processing/StreamingWorkScheduler.java:
##########
@@ -384,90 +361,210 @@ private ExecuteWorkResult executeWork(
stageInfo, computationState,
work.getLatencyTrackingId()));
try {
- WindmillStateReader stateReader = work.createWindmillStateReader();
- SideInputStateFetcher localSideInputStateFetcher =
-
sideInputStateFetcherFactory.createSideInputStateFetcher(work::fetchSideInput);
-
- // If the read output KVs, then we can decode Windmill's byte key into
userland
- // key object and provide it to the execution context for use with
per-key state.
- // Otherwise, we pass null.
- //
- // The coder type that will be present is:
- // WindowedValueCoder(TimerOrElementCoder(KvCoder))
- Optional<Coder<?>> keyCoder = computationWorkExecutor.keyCoder();
- @SuppressWarnings("deprecation")
- @Nullable
- final Object executionKey =
- !keyCoder.isPresent() ? null : keyCoder.get().decode(key.newInput(),
Coder.Context.OUTER);
-
- if (workItem.hasHotKeyInfo()) {
- Windmill.HotKeyInfo hotKeyInfo = workItem.getHotKeyInfo();
- Duration hotKeyAge = Duration.millis(hotKeyInfo.getHotKeyAgeUsec() /
1000);
-
- String stepName =
getShuffleTaskStepName(computationState.getMapTask());
- if (executionKey != null
- && (options.isHotKeyLoggingEnabled()
- || hasExperiment(options, "enable_hot_key_logging"))
- && keyCoder.isPresent()) {
- hotKeyLogger.logHotKeyDetection(stepName, hotKeyAge, executionKey);
- } else {
- hotKeyLogger.logHotKeyDetection(stepName, hotKeyAge);
- }
- }
+ StreamingModeExecutionContext context =
computationWorkExecutor.context();
// Blocks while executing work.
- computationWorkExecutor.executeWork(
- executionKey, work, stateReader, localSideInputStateFetcher,
outputBuilder);
+ computationWorkExecutor.executeWork(work, workExecutor, handle,
keyTransitionListener);
+
+ List<Work> workBatch;
+ List<Windmill.WorkItemCommitRequest> workItemCommits;
+ Map<Long, Pair<Instant, Runnable>> finalizationCallbacks;
+ long stateBytesRead;
+ {
+ if (context.workIsFailed()) {
+ throw new
WorkItemCancelledException(work.getWorkItem().getShardingKey());
+ }
+ context.flushState();
- if (work.isFailed()) {
- throw new WorkItemCancelledException(workItem.getShardingKey());
- }
+ // Retrieve executed works, work item commits, and accumulated
callbacks from execution
+ // context
+ workBatch = context.getExecutedWorks();
+ workItemCommits = context.getWorkItemCommits();
+ finalizationCallbacks = context.getFinalizationCallbacks();
+ stateBytesRead = context.getStateBytesRead();
- // Reports source bytes processed to WorkItemCommitRequest if available.
- try {
- long sourceBytesProcessed =
- computationWorkExecutor.computeSourceBytesProcessed(
- computationState.sourceBytesProcessCounterName());
- outputBuilder.setSourceBytesProcessed(sourceBytesProcessed);
- } catch (Exception e) {
- LOG.error("{}", e.toString());
+ context.reset(); // Don't use context after this.
}
-
-
commitFinalizer.cacheCommitFinalizers(computationWorkExecutor.context().flushState());
-
// Release the execution state for another thread to use.
computationState.releaseComputationWorkExecutor(computationWorkExecutor);
computationWorkExecutor = null;
- work.setState(Work.State.COMMIT_QUEUED);
-
outputBuilder.addAllPerWorkItemLatencyAttributions(work.getLatencyAttributions(sampler));
-
return ExecuteWorkResult.create(
- outputBuilder, stateReader.getBytesRead() +
localSideInputStateFetcher.getBytesRead());
+ workBatch, workItemCommits, finalizationCallbacks, stateBytesRead);
} catch (Throwable t) {
if (computationWorkExecutor != null) {
// If processing failed due to a thrown exception, close the
executionState. Do not
// return/release the executionState back to computationState as that
will lead to this
// executionState instance being reused.
- LOG.debug("Invalidating executor after work item {} failed",
workItem.getWorkToken(), t);
+ LOG.debug(
+ "Invalidating executor after work item {} failed",
+ work.getWorkItem().getWorkToken(),
+ t);
computationWorkExecutor.invalidate();
}
-
// Re-throw the exception, it will be caught and handled by
workFailureProcessor downstream.
throw t;
}
}
+ private void handleOnlyFinalize(
+ ComputationState computationState, Work work, Windmill.WorkItem
workItem) {
+ Windmill.WorkItemCommitRequest.Builder outputBuilder =
+ initializeOutputBuilder(workItem.getKey(), workItem);
+
outputBuilder.setSourceStateUpdates(Windmill.SourceState.newBuilder().setOnlyFinalize(true));
+ work.setState(Work.State.COMMIT_QUEUED);
+ work.queueCommit(outputBuilder.build(), computationState);
+ }
+
+ private StageInfo getStageInfo(ComputationState computationState) {
+ MapTask mapTask = computationState.getMapTask();
+ return stageInfoMap.computeIfAbsent(
+ mapTask.getStageName(), s -> StageInfo.create(s,
mapTask.getSystemName()));
+ }
+
+ private void commitWorkBatch(
+ ComputationState computationState,
+ List<Work> workBatch,
+ List<Windmill.WorkItemCommitRequest> workItemCommits) {
+ if (workBatch.isEmpty()) {
+ return;
+ }
+ if (workBatch.size() > 1 || multiKeyExperimentEnabled) {
+ commitMultiKeyWorkBatch(computationState, workBatch, workItemCommits);
+ } else {
+ commitSingleKeyWork(computationState, workBatch.get(0),
workItemCommits.get(0));
+ }
+ }
+
+ private void commitMultiKeyWorkBatch(
+ ComputationState computationState,
+ List<Work> workBatch,
+ List<Windmill.WorkItemCommitRequest> workItemCommits) {
+ Windmill.MultiKeyWorkItemCommitRequest.Builder multiKeyBuilder =
+ Windmill.MultiKeyWorkItemCommitRequest.newBuilder();
+
+ Work primaryWork = workBatch.get(0);
+ Work.KeyGroup keyGroup = primaryWork.getKeyGroup();
+ multiKeyBuilder.setKeyGroup(
+
Windmill.Uint128Proto.newBuilder().setHigh(keyGroup.high()).setLow(keyGroup.low()).build());
+
+ for (int i = 0; i < workBatch.size(); i++) {
+ // TODO: Add commit size validation
+ Windmill.WorkItemCommitRequest commit = workItemCommits.get(i);
Review Comment:
done.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
