scwhittle commented on code in PR #38919:
URL: https://github.com/apache/beam/pull/38919#discussion_r3632439486
##########
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:
should we just return nullable ExecutableWork instead? With annotations that
seems preferrable to me to avoid the optional allocation
##########
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:
if we don't use this often, perhaps we should just add the
ImmutableList.of() to the call-site
##########
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:
nit: I would uncapitalized E since reexecute is one word
##########
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:
could workBatch be an ImmutableList.Builder?
##########
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:
I don't see usage of this, can we just get rid of
this.getWorkStreamLatencies and instead do the recordGetWorkStreamLatencies
logic in the constructor?
##########
runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/StreamingModeExecutionContext.java:
##########
@@ -173,7 +225,13 @@ public StreamingModeExecutionContext(
StreamingModeExecutionStateRegistry executionStateRegistry,
StreamingGlobalConfigHandle globalConfigHandle,
long sinkByteLimit,
- boolean throwExceptionOnLargeOutput) {
+ boolean throwExceptionOnLargeOutput,
+ HotKeyLogger hotKeyLogger,
+ boolean hotKeyLoggingEnabled,
+ String stepName,
+ String sourceBytesProcessCounterName,
+ PipelineOptions options,
Review Comment:
instead of passing in full options, could pass some internal multikeyoptions
struct
benefits is that the parsing doesn't have to be here and can be shared
across all the contexts for now. But in the future we may want to configure
differently for different fused stages based upon other information and that
would allow us to do so as it would be separate from the single experiment
value.
##########
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:
precondition that workitemcommits and workBatch are same size
##########
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:
checkState that workBatch is non-empty
##########
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:
nit: can we use Cancelling with two ls to match other usage in beam
should we just use WorkItemCancelledException.java? If not can you explain
the differences here?
##########
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:
move activeWork definition above and use it here instead of separate check
##########
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:
can this instead be captured by looking at status? Maybe we can add a new
enum that indicates retryable abort separate from non-retryable abort?
or does this indicate some sort of rpc error where we didn't get a
commitstatus?
Add a comment.
--
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]
