diff --git a/CHANGES.md b/CHANGES.md index 59c7ac7b24ba..7ea2cccde299 100644 --- a/CHANGES.md +++ b/CHANGES.md @@ -70,6 +70,7 @@ * (Python) Removed the `envoy-data-plane` (and transitive `betterproto`) dependency; `EnvoyRateLimiter` now uses a small vendored protobuf definition instead, resolving dependency conflicts for downstream projects ([#37854](https://github.com/apache/beam/issues/37854)). * (Java) Supported acknowledge mode for JmsIO ([#39253](https://github.com/apache/beam/issues/39253)). +* (Python) Added `equal_to_approx`, an `assert_that` matcher that compares numeric pipeline outputs with a configurable tolerance ([#18028](https://github.com/apache/beam/issues/18028)). * X feature added (Java/Python) ([#X](https://github.com/apache/beam/issues/X)). ## Breaking Changes diff --git a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/windmill/client/WindmillStream.java b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/windmill/client/WindmillStream.java index 526b67890783..36001c151508 100644 --- a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/windmill/client/WindmillStream.java +++ b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/windmill/client/WindmillStream.java @@ -108,6 +108,11 @@ boolean commitWorkItem( Windmill.WorkItemCommitRequest request, Consumer onDone); + boolean commitMultiKeyWorkItem( + String computation, + Windmill.MultiKeyWorkItemCommitRequest request, + Consumer onDone); + /** Flushes any pending work items to the wire. */ void flush(); diff --git a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/windmill/client/commits/Commit.java b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/windmill/client/commits/Commit.java index b840d22a3434..bbd6cfc9432b 100644 --- a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/windmill/client/commits/Commit.java +++ b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/windmill/client/commits/Commit.java @@ -17,35 +17,89 @@ */ package org.apache.beam.runners.dataflow.worker.windmill.client.commits; -import com.google.auto.value.AutoValue; +import static org.apache.beam.sdk.util.Preconditions.checkStateNotNull; + import org.apache.beam.runners.dataflow.worker.streaming.ComputationState; import org.apache.beam.runners.dataflow.worker.streaming.Work; +import org.apache.beam.runners.dataflow.worker.windmill.Windmill.MultiKeyWorkItemCommitRequest; import org.apache.beam.runners.dataflow.worker.windmill.Windmill.WorkItemCommitRequest; import org.apache.beam.sdk.annotations.Internal; import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.Preconditions; +import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.ImmutableList; +import org.checkerframework.checker.nullness.qual.Nullable; /** Value class for a queued commit. */ @Internal -@AutoValue -public abstract class Commit { +public class Commit { + + private final ComputationState computationState; + private final ImmutableList workBatch; + private final @Nullable WorkItemCommitRequest singleKeyRequest; + private final @Nullable MultiKeyWorkItemCommitRequest multiKeyRequest; public static Commit create( WorkItemCommitRequest request, ComputationState computationState, Work work) { Preconditions.checkArgument(request.getSerializedSize() > 0); - return new AutoValue_Commit(request, computationState, work); + return new Commit(computationState, ImmutableList.of(work), request, null); + } + + public static Commit createMultiKey( + MultiKeyWorkItemCommitRequest multiKeyRequest, + ComputationState computationState, + ImmutableList workBatch) { + Preconditions.checkArgument(!workBatch.isEmpty()); + return new Commit(computationState, workBatch, null, multiKeyRequest); + } + + private Commit( + ComputationState computationState, + ImmutableList workBatch, + @Nullable WorkItemCommitRequest singleKeyRequest, + @Nullable MultiKeyWorkItemCommitRequest multiKeyRequest) { + this.computationState = computationState; + this.workBatch = workBatch; + this.singleKeyRequest = singleKeyRequest; + this.multiKeyRequest = multiKeyRequest; } public final String computationId() { return computationState().getComputationId(); } - public abstract WorkItemCommitRequest request(); + public @Nullable WorkItemCommitRequest singleKeyRequest() { + return singleKeyRequest; + }; - public abstract ComputationState computationState(); + public ComputationState computationState() { + return computationState; + } + + public @Nullable MultiKeyWorkItemCommitRequest multiKeyRequest() { + return multiKeyRequest; + } - public abstract Work work(); + public ImmutableList workBatch() { + return workBatch; + } + + public final int getSerializedByteSize() { + if (multiKeyRequest() != null) { + return checkStateNotNull(multiKeyRequest()).getSerializedSize(); + } + return checkStateNotNull(singleKeyRequest()).getSerializedSize(); + } - public final int getSize() { - return request().getSerializedSize(); + @Override + public String toString() { + Work work = workBatch.get(0); + return "[computationId=" + + computationId() + + ", shardingKey=" + + work.getShardedKey() + + ", workId=" + + work.id() + + ", workBatchSize=" + + workBatch.size() + + "]"; } } diff --git a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/windmill/client/commits/Commits.java b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/windmill/client/commits/Commits.java index 498e90f78e29..0607baebeb1d 100644 --- a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/windmill/client/commits/Commits.java +++ b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/windmill/client/commits/Commits.java @@ -31,6 +31,6 @@ public final class Commits { private Commits() {} public static WeightedSemaphore maxCommitByteSemaphore() { - return WeightedSemaphore.create(MAX_QUEUED_COMMITS_BYTES, Commit::getSize); + return WeightedSemaphore.create(MAX_QUEUED_COMMITS_BYTES, Commit::getSerializedByteSize); } } diff --git a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/windmill/client/commits/CompleteCommit.java b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/windmill/client/commits/CompleteCommit.java index e33e853d3d76..6c0a5a98e2ab 100644 --- a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/windmill/client/commits/CompleteCommit.java +++ b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/windmill/client/commits/CompleteCommit.java @@ -37,26 +37,11 @@ @AutoValue public abstract class CompleteCommit { - public static CompleteCommit create(Commit commit, CommitStatus commitStatus) { - return new AutoValue_CompleteCommit( - commit.computationId(), - ShardedKey.create(commit.request().getKey(), commit.request().getShardingKey()), - WorkId.builder() - .setWorkToken(commit.request().getWorkToken()) - .setCacheToken(commit.request().getCacheToken()) - .build(), - commitStatus); - } - public static CompleteCommit create( String computationId, ShardedKey shardedKey, WorkId workId, CommitStatus status) { return new AutoValue_CompleteCommit(computationId, shardedKey, workId, status); } - public static CompleteCommit forFailedWork(Commit commit) { - return create(commit, CommitStatus.ABORTED); - } - public abstract String computationId(); public abstract ShardedKey shardedKey(); diff --git a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/windmill/client/commits/StreamingApplianceWorkCommitter.java b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/windmill/client/commits/StreamingApplianceWorkCommitter.java index 20b95b0661d0..d627490dfe98 100644 --- a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/windmill/client/commits/StreamingApplianceWorkCommitter.java +++ b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/windmill/client/commits/StreamingApplianceWorkCommitter.java @@ -17,6 +17,9 @@ */ package org.apache.beam.runners.dataflow.worker.windmill.client.commits; +import static org.apache.beam.sdk.util.Preconditions.checkStateNotNull; +import static org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.Preconditions.checkState; + import java.util.HashMap; import java.util.Map; import java.util.concurrent.ExecutorService; @@ -112,7 +115,8 @@ private void commitLoop() { } while (commit != null) { ComputationState computationState = commit.computationState(); - commit.work().setState(Work.State.COMMITTING); + checkState(commit.workBatch().size() == 1); + commit.workBatch().get(0).setState(Work.State.COMMITTING); Windmill.ComputationCommitWorkRequest.Builder computationRequestBuilder = computationRequestMap.get(computationState); if (computationRequestBuilder == null) { @@ -120,10 +124,10 @@ private void commitLoop() { computationRequestBuilder.setComputationId(computationState.getComputationId()); computationRequestMap.put(computationState, computationRequestBuilder); } - computationRequestBuilder.addRequests(commit.request()); + computationRequestBuilder.addRequests(checkStateNotNull(commit.singleKeyRequest())); // Send the request if we've exceeded the bytes or there is no more // pending work. commitBytes is a long, so this cannot overflow. - commitBytes += commit.getSize(); + commitBytes += commit.getSerializedByteSize(); if (commitBytes >= TARGET_COMMIT_BUNDLE_BYTES) { break; } diff --git a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/windmill/client/commits/StreamingEngineWorkCommitter.java b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/windmill/client/commits/StreamingEngineWorkCommitter.java index b68f53121b86..83d0dfc6cda4 100644 --- a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/windmill/client/commits/StreamingEngineWorkCommitter.java +++ b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/windmill/client/commits/StreamingEngineWorkCommitter.java @@ -17,6 +17,8 @@ */ package org.apache.beam.runners.dataflow.worker.windmill.client.commits; +import static org.apache.beam.sdk.util.Preconditions.checkStateNotNull; + import com.google.auto.value.AutoBuilder; import java.util.concurrent.ExecutorService; import java.util.concurrent.Executors; @@ -30,6 +32,7 @@ import org.apache.beam.runners.dataflow.worker.streaming.WeightedBoundedQueue; import org.apache.beam.runners.dataflow.worker.streaming.WeightedSemaphore; import org.apache.beam.runners.dataflow.worker.streaming.Work; +import org.apache.beam.runners.dataflow.worker.windmill.Windmill.CommitStatus; import org.apache.beam.runners.dataflow.worker.windmill.client.CloseableStream; import org.apache.beam.runners.dataflow.worker.windmill.client.WindmillStream.CommitWorkStream; import org.apache.beam.sdk.annotations.Internal; @@ -100,8 +103,8 @@ public void start() { @Override public void commit(Commit commit) { - if (commit.work().isFailed()) { - failCommit(commit); + if (shouldFailCommit(commit)) { + failQueuedCommit(commit); } else { commitQueue.put(commit); } @@ -109,12 +112,7 @@ public void commit(Commit commit) { // Do this check after adding to commitQueue, else commitQueue.put() can race with // drainCommitQueue() in stop() and leave commits orphaned in the queue. if (!this.isRunning.get()) { - LOG.debug( - "Trying to queue commit on shutdown, failing commit=[computationId={}, shardingKey={}," - + " workId={} ].", - commit.computationId(), - commit.work().getShardedKey(), - commit.work().id()); + LOG.debug("Trying to queue commit on shutdown, failing commit={}", commit); drainCommitQueue(); } } @@ -141,14 +139,18 @@ public void stop() { private void drainCommitQueue() { Commit queuedCommit = commitQueue.poll(); while (queuedCommit != null) { - failCommit(queuedCommit); + failQueuedCommit(queuedCommit); queuedCommit = commitQueue.poll(); } } - private void failCommit(Commit commit) { - commit.work().setFailed(); - onCommitComplete.accept(CompleteCommit.forFailedWork(commit)); + private void failQueuedCommit(Commit commit) { + for (Work w : commit.workBatch()) { + w.setFailed(); + onCommitComplete.accept( + CompleteCommit.create( + commit.computationId(), w.getShardedKey(), w.id(), CommitStatus.ABORTED)); + } } @Override @@ -173,8 +175,8 @@ private void streamingCommitLoop() { // take() blocks until a value is available in the commitQueue. Preconditions.checkNotNull(initialCommit); - if (initialCommit.work().isFailed()) { - onCommitComplete.accept(CompleteCommit.forFailedWork(initialCommit)); + if (shouldFailCommit(initialCommit)) { + failQueuedCommit(initialCommit); initialCommit = null; continue; } @@ -194,29 +196,61 @@ private void streamingCommitLoop() { } } finally { if (initialCommit != null) { - failCommit(initialCommit); + failQueuedCommit(initialCommit); + } + } + } + + boolean shouldFailCommit(Commit commit) { + for (Work w : commit.workBatch()) { + if (w.isFailed()) { + return true; } } + return false; } /** Adds the commit to the batch if it fits, returning true if it is consumed. */ private boolean tryAddToCommitBatch(Commit commit, CommitWorkStream.RequestBatcher batcher) { Preconditions.checkNotNull(commit); - commit.work().setState(Work.State.COMMITTING); - activeCommitBytes.addAndGet(commit.getSize()); - boolean isCommitAccepted = - batcher.commitWorkItem( - commit.computationId(), - commit.request(), - commitStatus -> { - onCommitComplete.accept(CompleteCommit.create(commit, commitStatus)); - activeCommitBytes.addAndGet(-commit.getSize()); - }); + for (Work w : commit.workBatch()) { + w.setState(Work.State.COMMITTING); + } + activeCommitBytes.addAndGet(commit.getSerializedByteSize()); + boolean isCommitAccepted; + if (commit.multiKeyRequest() != null) { + isCommitAccepted = + batcher.commitMultiKeyWorkItem( + commit.computationId(), + checkStateNotNull(commit.multiKeyRequest()), + commitStatus -> { + for (Work w : commit.workBatch()) { + onCommitComplete.accept( + CompleteCommit.create( + commit.computationId(), w.getShardedKey(), w.id(), commitStatus)); + } + activeCommitBytes.addAndGet(-commit.getSerializedByteSize()); + }); + } else { + isCommitAccepted = + batcher.commitWorkItem( + commit.computationId(), + checkStateNotNull(commit.singleKeyRequest()), + commitStatus -> { + Work w = commit.workBatch().get(0); + onCommitComplete.accept( + CompleteCommit.create( + commit.computationId(), w.getShardedKey(), w.id(), commitStatus)); + activeCommitBytes.addAndGet(-commit.getSerializedByteSize()); + }); + } // Since the commit was not accepted, revert the changes made above. if (!isCommitAccepted) { - commit.work().setState(Work.State.COMMIT_QUEUED); - activeCommitBytes.addAndGet(-commit.getSize()); + for (Work w : commit.workBatch()) { + w.setState(Work.State.COMMIT_QUEUED); + } + activeCommitBytes.addAndGet(-commit.getSerializedByteSize()); } return isCommitAccepted; @@ -246,8 +280,8 @@ private boolean tryAddToCommitBatch(Commit commit, CommitWorkStream.RequestBatch } // Drop commits for failed work. Such commits will be dropped by Windmill anyway. - if (commit.work().isFailed()) { - onCommitComplete.accept(CompleteCommit.forFailedWork(commit)); + if (shouldFailCommit(commit)) { + failQueuedCommit(commit); continue; } diff --git a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/windmill/client/grpc/GrpcCommitWorkStream.java b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/windmill/client/grpc/GrpcCommitWorkStream.java index 160b0cce0133..e2a54b43cf07 100644 --- a/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/windmill/client/grpc/GrpcCommitWorkStream.java +++ b/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/windmill/client/grpc/GrpcCommitWorkStream.java @@ -34,6 +34,7 @@ import java.util.function.Function; import javax.annotation.Nullable; import org.apache.beam.repackaged.core.org.apache.commons.lang3.tuple.Pair; +import org.apache.beam.runners.dataflow.worker.windmill.Windmill; import org.apache.beam.runners.dataflow.worker.windmill.Windmill.CommitStatus; import org.apache.beam.runners.dataflow.worker.windmill.Windmill.JobHeader; import org.apache.beam.runners.dataflow.worker.windmill.Windmill.StreamingCommitRequestChunk; @@ -308,7 +309,7 @@ private void flushInternal(Map requests) if (requests.size() == 1) { Map.Entry elem = requests.entrySet().iterator().next(); - if (elem.getValue().getRequest().getSerializedSize() + if (elem.getValue().serializedCommit().size() > AbstractWindmillStream.RPC_STREAM_CHUNK_SIZE) { issueMultiChunkRequest(elem.getKey(), elem.getValue()); } else { @@ -324,9 +325,10 @@ private void issueSingleRequest(long id, PendingRequest pendingRequest) StreamingCommitWorkRequest.Builder requestBuilder = StreamingCommitWorkRequest.newBuilder(); requestBuilder .addCommitChunkBuilder() - .setComputationId(pendingRequest.getComputationId()) + .setComputationId(pendingRequest.computationId()) .setRequestId(id) .setShardingKey(pendingRequest.shardingKey()) + .setCommitType(pendingRequest.commitType()) .setSerializedWorkItemCommit(pendingRequest.serializedCommit()); StreamingCommitWorkRequest chunk = requestBuilder.build(); synchronized (this) { @@ -349,14 +351,15 @@ private void issueBatchedRequest(Map requests) for (Map.Entry entry : requests.entrySet()) { PendingRequest request = entry.getValue(); StreamingCommitRequestChunk.Builder chunkBuilder = requestBuilder.addCommitChunkBuilder(); - if (lastComputation == null || !lastComputation.equals(request.getComputationId())) { - chunkBuilder.setComputationId(request.getComputationId()); - lastComputation = request.getComputationId(); + if (lastComputation == null || !lastComputation.equals(request.computationId())) { + chunkBuilder.setComputationId(request.computationId()); + lastComputation = request.computationId(); } chunkBuilder .setRequestId(entry.getKey()) .setShardingKey(request.shardingKey()) - .setSerializedWorkItemCommit(request.serializedCommit()); + .setSerializedWorkItemCommit(request.serializedCommit()) + .setCommitType(request.commitType()); } StreamingCommitWorkRequest request = requestBuilder.build(); synchronized (this) { @@ -376,7 +379,7 @@ private void issueBatchedRequest(Map requests) private void issueMultiChunkRequest(long id, PendingRequest pendingRequest) throws WindmillStreamShutdownException { - checkNotNull(pendingRequest.getComputationId(), "Cannot commit WorkItem w/o a computationId."); + checkNotNull(pendingRequest.computationId(), "Cannot commit WorkItem w/o a computationId."); ByteString serializedCommit = pendingRequest.serializedCommit(); synchronized (this) { if (isShutdown) { @@ -397,8 +400,9 @@ private void issueMultiChunkRequest(long id, PendingRequest pendingRequest) StreamingCommitRequestChunk.newBuilder() .setRequestId(id) .setSerializedWorkItemCommit(chunk) - .setComputationId(pendingRequest.getComputationId()) - .setShardingKey(pendingRequest.shardingKey()); + .setComputationId(pendingRequest.computationId()) + .setShardingKey(pendingRequest.shardingKey()) + .setCommitType(pendingRequest.commitType()); int remaining = serializedCommit.size() - end; if (remaining > 0) { chunkBuilder.setRemainingBytesForWorkItem(remaining); @@ -416,24 +420,44 @@ private void issueMultiChunkRequest(long id, PendingRequest pendingRequest) private static class PendingRequest { private final String computationId; - private final WorkItemCommitRequest request; + private final long shardingKey; + private final ByteString serializedCommit; + private final StreamingCommitRequestChunk.CommitType commitType; private final Consumer onDone; private final long startTimeNanos; // System.nanoTime() of when request began. private PendingRequest( - String computationId, WorkItemCommitRequest request, Consumer onDone) { + String computationId, + long shardingKey, + ByteString serializedCommit, + StreamingCommitRequestChunk.CommitType commitType, + Consumer onDone) { this.computationId = computationId; - this.request = request; + this.shardingKey = shardingKey; + this.serializedCommit = serializedCommit; + this.commitType = commitType; this.onDone = onDone; this.startTimeNanos = System.nanoTime(); } - String getComputationId() { + String computationId() { return computationId; } - WorkItemCommitRequest getRequest() { - return request; + long shardingKey() { + return shardingKey; + } + + ByteString serializedCommit() { + return serializedCommit; + } + + StreamingCommitRequestChunk.CommitType commitType() { + return commitType; + } + + Consumer onDone() { + return onDone; } long getStartTimeNanos() { @@ -441,21 +465,13 @@ long getStartTimeNanos() { } private long getBytes() { - return (long) request.getSerializedSize() + computationId.length(); - } - - private ByteString serializedCommit() { - return request.toByteString(); + return (long) serializedCommit.size() + computationId.length(); } private void completeWithStatus(CommitStatus commitStatus) { onDone.accept(commitStatus); } - private long shardingKey() { - return request.getShardingKey(); - } - private void abort() { completeWithStatus(CommitStatus.ABORTED); } @@ -512,7 +528,34 @@ public boolean commitWorkItem( return false; } - PendingRequest request = new PendingRequest(computation, commitRequest, onDone); + PendingRequest request = + new PendingRequest( + computation, + commitRequest.getShardingKey(), + commitRequest.toByteString(), + StreamingCommitRequestChunk.CommitType.COMMIT_TYPE_SINGLE_KEY, + onDone); + add(idGenerator.incrementAndGet(), request); + return true; + } + + @Override + public boolean commitMultiKeyWorkItem( + String computation, + Windmill.MultiKeyWorkItemCommitRequest commitRequest, + Consumer onDone) { + Preconditions.checkArgument(commitRequest.getRequestsCount() > 0); + if (!canAccept(commitRequest.getSerializedSize() + computation.length())) { + return false; + } + PendingRequest request = + new PendingRequest( + computation, + // Any key in the batch for routing + commitRequest.getRequests(0).getShardingKey(), + commitRequest.toByteString(), + StreamingCommitRequestChunk.CommitType.COMMIT_TYPE_MULTI_KEY, + onDone); add(idGenerator.incrementAndGet(), request); return true; } diff --git a/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/FakeWindmillServer.java b/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/FakeWindmillServer.java index 5be8ec0a6c72..e5d68376a7d5 100644 --- a/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/FakeWindmillServer.java +++ b/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/FakeWindmillServer.java @@ -29,7 +29,6 @@ import java.util.ArrayList; import java.util.Collection; -import java.util.HashMap; import java.util.List; import java.util.Map; import java.util.Optional; @@ -37,10 +36,12 @@ import java.util.Set; import java.util.concurrent.ConcurrentHashMap; import java.util.concurrent.ConcurrentLinkedQueue; +import java.util.concurrent.CopyOnWriteArrayList; import java.util.concurrent.CountDownLatch; import java.util.concurrent.LinkedBlockingQueue; import java.util.concurrent.TimeUnit; import java.util.concurrent.atomic.AtomicInteger; +import java.util.concurrent.atomic.AtomicReference; import java.util.function.Consumer; import java.util.function.Function; import javax.annotation.concurrent.GuardedBy; @@ -49,6 +50,7 @@ import org.apache.beam.runners.dataflow.worker.streaming.WorkId; import org.apache.beam.runners.dataflow.worker.windmill.CloudWindmillMetadataServiceV1Alpha1Grpc; import org.apache.beam.runners.dataflow.worker.windmill.Windmill; +import org.apache.beam.runners.dataflow.worker.windmill.Windmill.CommitStatus; import org.apache.beam.runners.dataflow.worker.windmill.Windmill.CommitWorkResponse; import org.apache.beam.runners.dataflow.worker.windmill.Windmill.ComputationCommitWorkRequest; import org.apache.beam.runners.dataflow.worker.windmill.Windmill.ComputationGetDataRequest; @@ -87,8 +89,11 @@ public final class FakeWindmillServer extends WindmillServerStub { private final ResponseQueue dataToOffer; private final ResponseQueue commitsToOffer; private final Map streamingCommitsToOffer; + private final AtomicReference multiKeyCommitStatusToOffer; // Keys are work tokens. private final Map commitsReceived; + private final List multiKeyCommitsReceived = + new CopyOnWriteArrayList<>(); private final ArrayList statsReceived; private final LinkedBlockingQueue exceptions; private final AtomicInteger expectedExceptionCount; @@ -118,7 +123,9 @@ public FakeWindmillServer( commitsToOffer = new ResponseQueue() .returnByDefault(CommitWorkResponse.getDefaultInstance()); - streamingCommitsToOffer = new HashMap<>(); + streamingCommitsToOffer = new ConcurrentHashMap<>(); + // Respond multikey commits with ok, unless overridden. + multiKeyCommitStatusToOffer = new AtomicReference<>(CommitStatus.OK); commitsReceived = new ConcurrentHashMap<>(); exceptions = new LinkedBlockingQueue<>(); expectedExceptionCount = new AtomicInteger(); @@ -153,6 +160,11 @@ public Map whenCommitWorkStreamCalled() { return streamingCommitsToOffer; } + /** @param commitStatus status to return to multiKeyCommits */ + public void setMultiKeyCommitStatus(CommitStatus commitStatus) { + this.multiKeyCommitStatusToOffer.set(commitStatus); + } + @Override public Windmill.GetWorkResponse getWork(Windmill.GetWorkRequest request) { LOG.debug("getWorkRequest: {}", request.toString()); @@ -400,6 +412,7 @@ public void shutdown() {} public RequestBatcher batcher() { return new RequestBatcher() { final List requests = new ArrayList<>(); + final List multiKeyRequests = new ArrayList<>(); @Override public boolean commitWorkItem( @@ -423,6 +436,17 @@ public boolean commitWorkItem( return true; } + @Override + public boolean commitMultiKeyWorkItem( + String computation, + Windmill.MultiKeyWorkItemCommitRequest request, + Consumer onDone) { + LOG.debug("commitWorkStream::commitMultiKeyWorkItem: {}", request); + multiKeyRequests.add(new MultiKeyRequestAndDone(request, onDone)); + flush(); + return true; + } + @Override public void flush() { for (RequestAndDone elem : requests) { @@ -445,6 +469,24 @@ public void flush() { .orElse(Windmill.CommitStatus.OK)); } requests.clear(); + + for (MultiKeyRequestAndDone elem : multiKeyRequests) { + if (dropStreamingCommits) { + for (WorkItemCommitRequest workRequest : elem.request.getRequestsList()) { + droppedStreamingCommits.put(workRequest.getWorkToken(), elem.onDone); + } + continue; + } + + multiKeyCommitsReceived.add(elem.request); + for (WorkItemCommitRequest workRequest : elem.request.getRequestsList()) { + commitsReceived.put(workRequest.getWorkToken(), workRequest); + } + + Windmill.CommitStatus status = multiKeyCommitStatusToOffer.get(); + elem.onDone.accept(status); + } + multiKeyRequests.clear(); } class RequestAndDone { @@ -456,6 +498,18 @@ class RequestAndDone { this.onDone = onDone; } } + + class MultiKeyRequestAndDone { + final Consumer onDone; + final Windmill.MultiKeyWorkItemCommitRequest request; + + MultiKeyRequestAndDone( + Windmill.MultiKeyWorkItemCommitRequest request, + Consumer onDone) { + this.request = request; + this.onDone = onDone; + } + } }; } @@ -518,6 +572,15 @@ public Map waitForAndGetCommits(int numCommits) { public void clearCommitsReceived() { commitsRequested = 0; commitsReceived.clear(); + multiKeyCommitsReceived.clear(); + } + + public List getMultiKeyCommitsReceived() { + return multiKeyCommitsReceived; + } + + public void clearMultiKeyCommitsReceived() { + multiKeyCommitsReceived.clear(); } public ConcurrentHashMap> waitForDroppedCommits( diff --git a/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/windmill/client/commits/StreamingApplianceWorkCommitterTest.java b/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/windmill/client/commits/StreamingApplianceWorkCommitterTest.java index 5c3132ae471d..0596210a0270 100644 --- a/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/windmill/client/commits/StreamingApplianceWorkCommitterTest.java +++ b/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/windmill/client/commits/StreamingApplianceWorkCommitterTest.java @@ -128,10 +128,11 @@ public void testCommit() { fakeWindmillServer.waitForAndGetCommits(commits.size()); for (Commit commit : commits) { + assertThat(commit.workBatch()).hasSize(1); Windmill.WorkItemCommitRequest request = - committed.get(commit.work().getWorkItem().getWorkToken()); + committed.get(commit.workBatch().get(0).getWorkItem().getWorkToken()); assertNotNull(request); - assertThat(request).isEqualTo(commit.request()); + assertThat(request).isEqualTo(commit.singleKeyRequest()); } assertThat(completeCommits).hasSize(commits.size()); @@ -141,12 +142,14 @@ public void testCommit() { (CompleteCommit completeCommit, Commit commit) -> completeCommit.computationId().equals(commit.computationId()) && completeCommit.status() == Windmill.CommitStatus.OK - && completeCommit.workId().equals(commit.work().id()) + && commit.workBatch().size() == 1 + && completeCommit.workId().equals(commit.workBatch().get(0).id()) && completeCommit .shardedKey() .equals( ShardedKey.create( - commit.request().getKey(), commit.request().getShardingKey())), + commit.singleKeyRequest().getKey(), + commit.singleKeyRequest().getShardingKey())), "expected to equal")) .containsExactlyElementsIn(commits); } diff --git a/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/windmill/client/commits/StreamingEngineWorkCommitterTest.java b/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/windmill/client/commits/StreamingEngineWorkCommitterTest.java index 01197622c24d..881cf620e8d8 100644 --- a/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/windmill/client/commits/StreamingEngineWorkCommitterTest.java +++ b/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/windmill/client/commits/StreamingEngineWorkCommitterTest.java @@ -19,6 +19,7 @@ import static com.google.common.truth.Truth.assertThat; import static org.apache.beam.runners.dataflow.worker.windmill.Windmill.CommitStatus.OK; +import static org.junit.Assert.assertEquals; import static org.junit.Assert.assertNotNull; import static org.junit.Assert.assertTrue; import static org.mockito.Mockito.mock; @@ -53,6 +54,7 @@ import org.apache.beam.runners.dataflow.worker.streaming.WorkId; import org.apache.beam.runners.dataflow.worker.util.BoundedQueueExecutor; import org.apache.beam.runners.dataflow.worker.windmill.Windmill; +import org.apache.beam.runners.dataflow.worker.windmill.Windmill.CommitStatus; import org.apache.beam.runners.dataflow.worker.windmill.Windmill.WorkItem; import org.apache.beam.runners.dataflow.worker.windmill.Windmill.WorkItemCommitRequest; import org.apache.beam.runners.dataflow.worker.windmill.client.CloseableStream; @@ -62,6 +64,7 @@ import org.apache.beam.runners.dataflow.worker.windmill.work.refresh.HeartbeatSender; import org.apache.beam.vendor.grpc.v1p69p0.com.google.protobuf.ByteString; import org.apache.beam.vendor.grpc.v1p69p0.io.grpc.testing.GrpcCleanupRule; +import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.ImmutableList; import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.ImmutableMap; import org.joda.time.Duration; import org.joda.time.Instant; @@ -134,12 +137,10 @@ private static ComputationState createComputationState(String computationId) { null); } - private static CompleteCommit asCompleteCommit(Commit commit, Windmill.CommitStatus status) { - if (commit.work().isFailed()) { - return CompleteCommit.forFailedWork(commit); - } - - return CompleteCommit.create(commit, status); + private static CompleteCommit asCompleteCommit( + String computationId, Work work, Windmill.CommitStatus status) { + Windmill.CommitStatus finalStatus = work.isFailed() ? Windmill.CommitStatus.ABORTED : status; + return CompleteCommit.create(computationId, work.getShardedKey(), work.id(), finalStatus); } @Before @@ -186,10 +187,15 @@ public void testCommit_sendsCommitsToStreamingEngine() { waitForExpectedSetSize(completeCommits, 5); for (Commit commit : commits) { - WorkItemCommitRequest request = committed.get(commit.work().getWorkItem().getWorkToken()); + assertThat(commit.workBatch()).hasSize(1); + WorkItemCommitRequest request = + committed.get(commit.workBatch().get(0).getWorkItem().getWorkToken()); assertNotNull(request); - assertThat(request).isEqualTo(commit.request()); - assertThat(completeCommits).contains(asCompleteCommit(commit, Windmill.CommitStatus.OK)); + assertThat(request).isEqualTo(commit.singleKeyRequest()); + assertThat(completeCommits) + .contains( + asCompleteCommit( + commit.computationId(), commit.workBatch().get(0), Windmill.CommitStatus.OK)); } workCommitter.stop(); @@ -224,14 +230,24 @@ public void testCommit_handlesFailedCommits() { waitForExpectedSetSize(completeCommits, 10); for (Commit commit : commits) { - if (commit.work().isFailed()) { + assertThat(commit.workBatch()).hasSize(1); + if (commit.workBatch().get(0).isFailed()) { assertThat(completeCommits) - .contains(asCompleteCommit(commit, Windmill.CommitStatus.ABORTED)); - assertThat(committed).doesNotContainKey(commit.work().getWorkItem().getWorkToken()); + .contains( + asCompleteCommit( + commit.computationId(), + commit.workBatch().get(0), + Windmill.CommitStatus.ABORTED)); + assertThat(committed) + .doesNotContainKey(commit.workBatch().get(0).getWorkItem().getWorkToken()); } else { - assertThat(completeCommits).contains(asCompleteCommit(commit, Windmill.CommitStatus.OK)); + assertThat(completeCommits) + .contains( + asCompleteCommit( + commit.computationId(), commit.workBatch().get(0), Windmill.CommitStatus.OK)); assertThat(committed) - .containsEntry(commit.work().getWorkItem().getWorkToken(), commit.request()); + .containsEntry( + commit.workBatch().get(0).getWorkItem().getWorkToken(), commit.singleKeyRequest()); } } @@ -282,11 +298,17 @@ public void testCommit_handlesCompleteCommits_commitStatusNotOK() { waitForExpectedSetSize(completeCommits, commits.size()); for (Commit commit : commits) { - WorkItemCommitRequest request = committed.get(commit.work().getWorkItem().getWorkToken()); + assertThat(commit.workBatch()).hasSize(1); + WorkItemCommitRequest request = + committed.get(commit.workBatch().get(0).getWorkItem().getWorkToken()); assertNotNull(request); - assertThat(request).isEqualTo(commit.request()); + assertThat(request).isEqualTo(commit.singleKeyRequest()); assertThat(completeCommits) - .contains(asCompleteCommit(commit, expectedCommitStatus.get(commit.work().id()))); + .contains( + asCompleteCommit( + commit.computationId(), + commit.workBatch().get(0), + expectedCommitStatus.get(commit.workBatch().get(0).id()))); } workCommitter.stop(); @@ -313,6 +335,14 @@ public boolean commitWorkItem( return false; } + @Override + public boolean commitMultiKeyWorkItem( + String computation, + Windmill.MultiKeyWorkItemCommitRequest request, + Consumer onDone) { + return false; + } + @Override public void flush() {} }; @@ -370,7 +400,8 @@ public void shutdown() {} } for (Commit commit : commits) { - assertTrue(commit.work().isFailed()); + assertThat(commit.workBatch()).hasSize(1); + assertTrue(commit.workBatch().get(0).isFailed()); } } @@ -409,10 +440,15 @@ public void testMultipleCommitSendersSingleStream() { waitForExpectedSetSize(completeCommits, commits.size()); for (Commit commit : commits) { - WorkItemCommitRequest request = committed.get(commit.work().getWorkItem().getWorkToken()); + assertThat(commit.workBatch()).hasSize(1); + WorkItemCommitRequest request = + committed.get(commit.workBatch().get(0).getWorkItem().getWorkToken()); assertNotNull(request); - assertThat(request).isEqualTo(commit.request()); - assertThat(completeCommits).contains(asCompleteCommit(commit, Windmill.CommitStatus.OK)); + assertThat(request).isEqualTo(commit.singleKeyRequest()); + assertThat(completeCommits) + .contains( + asCompleteCommit( + commit.computationId(), commit.workBatch().get(0), Windmill.CommitStatus.OK)); } workCommitter.stop(); @@ -474,4 +510,207 @@ public void testStop_drainsCommitQueue_concurrentCommit() waitForExpectedSetSize(completeCommits, sentCommits.intValue()); } + + @Test + public void testCommit_multiKeyCommitSuccess() { + Set completeCommits = Collections.newSetFromMap(new ConcurrentHashMap<>()); + workCommitter = createWorkCommitter(completeCommits::add); + + Work workA = createMockWork(101L); + Work workB = createMockWork(102L); + Work workC = createMockWork(103L); + + Windmill.MultiKeyWorkItemCommitRequest multiKeyRequest = + Windmill.MultiKeyWorkItemCommitRequest.newBuilder() + .addRequests( + Windmill.WorkItemCommitRequest.newBuilder() + .setKey(workA.getWorkItem().getKey()) + .setShardingKey(workA.getWorkItem().getShardingKey()) + .setWorkToken(workA.getWorkItem().getWorkToken()) + .setCacheToken(workA.getWorkItem().getCacheToken()) + .build()) + .addRequests( + Windmill.WorkItemCommitRequest.newBuilder() + .setKey(workB.getWorkItem().getKey()) + .setShardingKey(workB.getWorkItem().getShardingKey()) + .setWorkToken(workB.getWorkItem().getWorkToken()) + .setCacheToken(workB.getWorkItem().getCacheToken()) + .build()) + .addRequests( + Windmill.WorkItemCommitRequest.newBuilder() + .setKey(workC.getWorkItem().getKey()) + .setShardingKey(workC.getWorkItem().getShardingKey()) + .setWorkToken(workC.getWorkItem().getWorkToken()) + .setCacheToken(workC.getWorkItem().getCacheToken()) + .build()) + .build(); + + Commit commit = + Commit.createMultiKey( + multiKeyRequest, + createComputationState("computationId"), + ImmutableList.of(workA, workB, workC)); + + workCommitter.start(); + workCommitter.commit(commit); + + // Wait for the server to receive and process the commits + fakeWindmillServer.waitForAndGetCommits(3); + waitForExpectedSetSize(completeCommits, 3); + + // Verify that FakeWindmillServer received all 3 work requests in multiKeyCommitsReceived + List multiKeyCommits = + fakeWindmillServer.getMultiKeyCommitsReceived(); + assertThat(multiKeyCommits).hasSize(1); + assertThat(multiKeyCommits.get(0)).isEqualTo(multiKeyRequest); + + // Verify all three works are completed successfully + assertThat(completeCommits) + .containsExactly( + CompleteCommit.create( + "computationId", workA.getShardedKey(), workA.id(), CommitStatus.OK), + CompleteCommit.create( + "computationId", workB.getShardedKey(), workB.id(), CommitStatus.OK), + CompleteCommit.create( + "computationId", workC.getShardedKey(), workC.id(), CommitStatus.OK)); + + // There should be no more commits in the queue + assertEquals(0, workCommitter.currentActiveCommitBytes()); + workCommitter.stop(); + } + + @Test + public void testCommit_multiKeyCommitFailedWork() { + Set completeCommits = Collections.newSetFromMap(new ConcurrentHashMap<>()); + workCommitter = createWorkCommitter(completeCommits::add); + + Work workA = createMockWork(101L); + Work workB = createMockWork(102L); + Work workC = createMockWork(103L); + + // Mark non-primary key B as failed + workB.setFailed(); + + Windmill.MultiKeyWorkItemCommitRequest multiKeyRequest = + Windmill.MultiKeyWorkItemCommitRequest.newBuilder() + .addRequests( + Windmill.WorkItemCommitRequest.newBuilder() + .setKey(workA.getWorkItem().getKey()) + .setShardingKey(workA.getWorkItem().getShardingKey()) + .setWorkToken(workA.getWorkItem().getWorkToken()) + .setCacheToken(workA.getWorkItem().getCacheToken()) + .build()) + .addRequests( + Windmill.WorkItemCommitRequest.newBuilder() + .setKey(workB.getWorkItem().getKey()) + .setShardingKey(workB.getWorkItem().getShardingKey()) + .setWorkToken(workB.getWorkItem().getWorkToken()) + .setCacheToken(workB.getWorkItem().getCacheToken()) + .build()) + .addRequests( + Windmill.WorkItemCommitRequest.newBuilder() + .setKey(workC.getWorkItem().getKey()) + .setShardingKey(workC.getWorkItem().getShardingKey()) + .setWorkToken(workC.getWorkItem().getWorkToken()) + .setCacheToken(workC.getWorkItem().getCacheToken()) + .build()) + .build(); + + Commit commit = + Commit.createMultiKey( + multiKeyRequest, + createComputationState("computationId"), + ImmutableList.of(workA, workB, workC)); + + workCommitter.start(); + workCommitter.commit(commit); + + // The entire batch must be aborted immediately without making network calls + waitForExpectedSetSize(completeCommits, 3); + + // Verify all three works are aborted individually + assertThat(completeCommits) + .containsExactly( + CompleteCommit.create( + "computationId", workA.getShardedKey(), workA.id(), CommitStatus.ABORTED), + CompleteCommit.create( + "computationId", workB.getShardedKey(), workB.id(), CommitStatus.ABORTED), + CompleteCommit.create( + "computationId", workC.getShardedKey(), workC.id(), CommitStatus.ABORTED)); + + // There should be no more commits in the queue + assertEquals(0, workCommitter.currentActiveCommitBytes()); + workCommitter.stop(); + } + + @Test + public void testCommit_multiKeyCommitStatusNotOK() { + Set completeCommits = Collections.newSetFromMap(new ConcurrentHashMap<>()); + workCommitter = createWorkCommitter(completeCommits::add); + + Work workA = createMockWork(101L); + Work workB = createMockWork(102L); + Work workC = createMockWork(103L); + + Windmill.MultiKeyWorkItemCommitRequest multiKeyRequest = + Windmill.MultiKeyWorkItemCommitRequest.newBuilder() + .addRequests( + Windmill.WorkItemCommitRequest.newBuilder() + .setKey(workA.getWorkItem().getKey()) + .setShardingKey(workA.getWorkItem().getShardingKey()) + .setWorkToken(workA.getWorkItem().getWorkToken()) + .setCacheToken(workA.getWorkItem().getCacheToken()) + .build()) + .addRequests( + Windmill.WorkItemCommitRequest.newBuilder() + .setKey(workB.getWorkItem().getKey()) + .setShardingKey(workB.getWorkItem().getShardingKey()) + .setWorkToken(workB.getWorkItem().getWorkToken()) + .setCacheToken(workB.getWorkItem().getCacheToken()) + .build()) + .addRequests( + Windmill.WorkItemCommitRequest.newBuilder() + .setKey(workC.getWorkItem().getKey()) + .setShardingKey(workC.getWorkItem().getShardingKey()) + .setWorkToken(workC.getWorkItem().getWorkToken()) + .setCacheToken(workC.getWorkItem().getCacheToken()) + .build()) + .build(); + + Commit commit = + Commit.createMultiKey( + multiKeyRequest, + createComputationState("computationId"), + ImmutableList.of(workA, workB, workC)); + + // Respond to multi key commit with NOT_FOUND status. + fakeWindmillServer.setMultiKeyCommitStatus(CommitStatus.NOT_FOUND); + + workCommitter.start(); + workCommitter.commit(commit); + + // Wait for the server to receive and process the commits + fakeWindmillServer.waitForAndGetCommits(3); + waitForExpectedSetSize(completeCommits, 3); + + // Verify that FakeWindmillServer received the multi-key commit + List multiKeyCommits = + fakeWindmillServer.getMultiKeyCommitsReceived(); + assertThat(multiKeyCommits).hasSize(1); + assertThat(multiKeyCommits.get(0)).isEqualTo(multiKeyRequest); + + // Verify all three works in the multi-key commit are completed with NOT_FOUND status + assertThat(completeCommits) + .containsExactly( + CompleteCommit.create( + "computationId", workA.getShardedKey(), workA.id(), CommitStatus.NOT_FOUND), + CompleteCommit.create( + "computationId", workB.getShardedKey(), workB.id(), CommitStatus.NOT_FOUND), + CompleteCommit.create( + "computationId", workC.getShardedKey(), workC.id(), CommitStatus.NOT_FOUND)); + + // There should be no more commits in the queue + assertEquals(0, workCommitter.currentActiveCommitBytes()); + workCommitter.stop(); + } } diff --git a/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/windmill/client/grpc/GrpcCommitWorkStreamTest.java b/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/windmill/client/grpc/GrpcCommitWorkStreamTest.java index 9c3d5c9c3ef3..1e995f4047c3 100644 --- a/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/windmill/client/grpc/GrpcCommitWorkStreamTest.java +++ b/runners/google-cloud-dataflow-java/worker/src/test/java/org/apache/beam/runners/dataflow/worker/windmill/client/grpc/GrpcCommitWorkStreamTest.java @@ -42,6 +42,8 @@ import java.util.function.Supplier; import org.apache.beam.runners.dataflow.worker.windmill.CloudWindmillServiceV1Alpha1Grpc; import org.apache.beam.runners.dataflow.worker.windmill.Windmill; +import org.apache.beam.runners.dataflow.worker.windmill.Windmill.CommitStatus; +import org.apache.beam.runners.dataflow.worker.windmill.Windmill.StreamingCommitResponse; import org.apache.beam.runners.dataflow.worker.windmill.WindmillConnection; import org.apache.beam.runners.dataflow.worker.windmill.client.TriggeredScheduledExecutorService; import org.apache.beam.runners.dataflow.worker.windmill.client.WindmillStream; @@ -1134,6 +1136,89 @@ public void testCommitWorkItem_multiplePhysicalStreams_multipleHandovers_halfClo assertTrue(commitWorkStream.awaitTermination(10, TimeUnit.SECONDS)); } + @Test + public void testCommit_multiKeyCommit() throws Exception { + testMultiKeyCommit(CommitStatus.OK); + } + + @Test + public void testCommit_multiKeyCommit_Failure() throws Exception { + testMultiKeyCommit(CommitStatus.NOT_FOUND); + } + + private void testMultiKeyCommit(CommitStatus commitStatus) throws Exception { + GrpcCommitWorkStream commitWorkStream = createCommitWorkStream(); + FakeWindmillGrpcService.CommitStreamInfo streamInfo = waitForConnectionAndConsumeHeader(); + + CompletableFuture commitStatusFuture = new CompletableFuture<>(); + + // 1. Construct two individual WorkItemCommitRequests + long shardingKey1 = 101L; + long workToken1 = 201L; + long cacheToken1 = 301L; + long shardingKey2 = 102L; + long workToken2 = 202L; + long cacheToken2 = 302L; + Windmill.WorkItemCommitRequest request1 = + Windmill.WorkItemCommitRequest.newBuilder() + .setKey(ByteString.copyFromUtf8("key1")) + .setShardingKey(shardingKey1) + .setWorkToken(workToken1) + .setCacheToken(cacheToken1) + .build(); + Windmill.WorkItemCommitRequest request2 = + Windmill.WorkItemCommitRequest.newBuilder() + .setKey(ByteString.copyFromUtf8("key2")) + .setShardingKey(shardingKey2) + .setWorkToken(workToken2) + .setCacheToken(cacheToken2) + .build(); + + // 2. Wrap them into a MultiKeyWorkItemCommitRequest + Windmill.MultiKeyWorkItemCommitRequest multiKeyRequest = + Windmill.MultiKeyWorkItemCommitRequest.newBuilder() + .addRequests(request1) + .addRequests(request2) + .build(); + + // 3. Commit the multi-key work item using the request batcher + try (WindmillStream.CommitWorkStream.RequestBatcher batcher = commitWorkStream.batcher()) { + assertTrue( + batcher.commitMultiKeyWorkItem( + COMPUTATION_ID, multiKeyRequest, commitStatusFuture::complete)); + } + + // 4. Receive and assert request properties on FakeWindmillGrpcService + Windmill.StreamingCommitWorkRequest request = streamInfo.requests.take(); + assertThat(request.getCommitChunkCount()).isEqualTo(1); + + Windmill.StreamingCommitRequestChunk chunk = request.getCommitChunk(0); + + // Assert that the commit type is correctly identified as COMMIT_TYPE_MULTI_KEY + assertThat(chunk.getCommitType()) + .isEqualTo(Windmill.StreamingCommitRequestChunk.CommitType.COMMIT_TYPE_MULTI_KEY); + + // Assert that the routing sharding key is mapped to the first request's sharding key + assertThat(chunk.getShardingKey()).isEqualTo(request1.getShardingKey()); + + // Assert that the serialized payload matches the input multiKeyRequest + Windmill.MultiKeyWorkItemCommitRequest parsedRequest = + Windmill.MultiKeyWorkItemCommitRequest.parseFrom(chunk.getSerializedWorkItemCommit()); + assertThat(parsedRequest).isEqualTo(multiKeyRequest); + + // 5. Respond with the generated requestId to complete the commit + long requestId = chunk.getRequestId(); + StreamingCommitResponse.Builder builder = + StreamingCommitResponse.newBuilder().addRequestId(requestId); + if (commitStatus != CommitStatus.OK) { + builder.addStatus(commitStatus); + } + streamInfo.responseObserver.onNext(builder.build()); + + // 6. Verify callback completed with expected sCommitStatus + assertThat(commitStatusFuture.get()).isEqualTo(commitStatus); + } + @Test public void testCommitWorkItem_stopsRetriesAfterDuration() throws Exception { int numCommits = 1; diff --git a/runners/google-cloud-dataflow-java/worker/windmill/src/main/proto/windmill.proto b/runners/google-cloud-dataflow-java/worker/windmill/src/main/proto/windmill.proto index aaa09c105fc3..a7a99e2ca5a1 100644 --- a/runners/google-cloud-dataflow-java/worker/windmill/src/main/proto/windmill.proto +++ b/runners/google-cloud-dataflow-java/worker/windmill/src/main/proto/windmill.proto @@ -678,9 +678,24 @@ message WorkItemCommitRequest { reserved 6, 23; } +message MultiKeyWorkItemCommitRequest { + optional Uint128Proto key_group = 7; + + repeated WorkItemCommitRequest requests = 1; + + repeated OutputMessageBundle output_messages = 2; + + repeated PubSubMessageBundle pubsub_messages = 3; + + repeated int64 finalize_ids = 4 [packed = true]; + + reserved 6; +} + message ComputationCommitWorkRequest { required string computation_id = 1; repeated WorkItemCommitRequest requests = 2; + repeated MultiKeyWorkItemCommitRequest multi_key_requests = 3; } message CommitWorkRequest { @@ -906,6 +921,14 @@ message StreamingCommitRequestChunk { // before handing off to the WindmillHost for processing. optional int64 remaining_bytes_for_work_item = 4; optional bytes serialized_work_item_commit = 5; + + enum CommitType { + COMMIT_TYPE_UNSPECIFIED = 0; + COMMIT_TYPE_SINGLE_KEY = 1; + COMMIT_TYPE_MULTI_KEY = 2; + } + + optional CommitType commit_type = 7; } message StreamingCommitResponse { diff --git a/sdks/java/extensions/google-cloud-platform-core/build.gradle b/sdks/java/extensions/google-cloud-platform-core/build.gradle index f1bfb63c7a3a..78cfe4739ec5 100644 --- a/sdks/java/extensions/google-cloud-platform-core/build.gradle +++ b/sdks/java/extensions/google-cloud-platform-core/build.gradle @@ -56,6 +56,7 @@ dependencies { implementation library.java.http_core implementation library.java.http_client implementation library.java.jackson_annotations + implementation library.java.jackson_core implementation library.java.jackson_databind permitUnusedDeclared library.java.jackson_databind // BEAM-11761 testImplementation project(path: ":sdks:java:core", configuration: "shadowTest") diff --git a/sdks/java/extensions/google-cloud-platform-core/src/main/java/org/apache/beam/sdk/extensions/gcp/options/GcsOptions.java b/sdks/java/extensions/google-cloud-platform-core/src/main/java/org/apache/beam/sdk/extensions/gcp/options/GcsOptions.java index 2da382a5b674..134c4cb3f281 100644 --- a/sdks/java/extensions/google-cloud-platform-core/src/main/java/org/apache/beam/sdk/extensions/gcp/options/GcsOptions.java +++ b/sdks/java/extensions/google-cloud-platform-core/src/main/java/org/apache/beam/sdk/extensions/gcp/options/GcsOptions.java @@ -18,6 +18,15 @@ package org.apache.beam.sdk.extensions.gcp.options; import com.fasterxml.jackson.annotation.JsonIgnore; +import com.fasterxml.jackson.core.JsonGenerator; +import com.fasterxml.jackson.core.JsonParser; +import com.fasterxml.jackson.databind.DeserializationContext; +import com.fasterxml.jackson.databind.JsonDeserializer; +import com.fasterxml.jackson.databind.JsonNode; +import com.fasterxml.jackson.databind.JsonSerializer; +import com.fasterxml.jackson.databind.SerializerProvider; +import com.fasterxml.jackson.databind.annotation.JsonDeserialize; +import com.fasterxml.jackson.databind.annotation.JsonSerialize; import com.google.cloud.hadoop.gcsio.GoogleCloudStorageReadOptions; import com.google.cloud.hadoop.util.AsyncWriteChannelOptions; import java.util.HashMap; @@ -49,12 +58,16 @@ public interface GcsOptions extends ApplicationNameOptions, GcpOptions, Pipeline class GcsReadOptionsFactory implements DefaultValueFactory { @Override public GoogleCloudStorageReadOptions create(PipelineOptions options) { - return GoogleCloudStorageReadOptions.DEFAULT; + // In gcs-connector v3, GoogleCloudStorageReadOptions.DEFAULT changed fadvise from SEQUENTIAL + // to AUTO. Beam workloads default to SEQUENTIAL to preserve expected sequential read + // throughput and caching behavior. + return GcsReadOptionsSerializer.DEFAULT_OPTIONS; } } /** @deprecated This option will be removed in a future release. */ - @JsonIgnore + @JsonSerialize(using = GcsReadOptionsSerializer.class) + @JsonDeserialize(using = GcsReadOptionsDeserializer.class) @Description( "The GoogleCloudStorageReadOptions instance that should be used to read from Google Cloud Storage.") @Default.InstanceFactory(GcsReadOptionsFactory.class) @@ -286,3 +299,71 @@ boolean exceedsEntryLimit() { } } } + +class GcsReadOptionsSerializer extends JsonSerializer { + static final GoogleCloudStorageReadOptions DEFAULT_OPTIONS = + GoogleCloudStorageReadOptions.DEFAULT + .toBuilder() + .setFadvise(GoogleCloudStorageReadOptions.Fadvise.SEQUENTIAL) + .build(); + + @Override + public void serialize( + GoogleCloudStorageReadOptions value, JsonGenerator gen, SerializerProvider serializers) + throws java.io.IOException { + // Note: We only support a partial set of options to propagate to remote + // workers. Setting the unsupported ones will not have any effect and will fall + // back to default in remote workers. Support for additional options can be + // added on a need basis. The full list of options can be seen in + // com.google.cloud.hadoop.gcsio.GoogleCloudStorageReadOptions. + gen.writeStartObject(); + if (value.getFadvise() != null && value.getFadvise() != DEFAULT_OPTIONS.getFadvise()) { + gen.writeStringField("fadvise", value.getFadvise().name()); + } + if (value.isFastFailOnNotFoundEnabled() != DEFAULT_OPTIONS.isFastFailOnNotFoundEnabled()) { + gen.writeBooleanField("fastFailOnNotFoundEnabled", value.isFastFailOnNotFoundEnabled()); + } + if (value.getMinRangeRequestSize() != DEFAULT_OPTIONS.getMinRangeRequestSize()) { + gen.writeNumberField("minRangeRequestSize", value.getMinRangeRequestSize()); + } + if (value.getInplaceSeekLimit() != DEFAULT_OPTIONS.getInplaceSeekLimit()) { + gen.writeNumberField("inplaceSeekLimit", value.getInplaceSeekLimit()); + } + if (value.isGrpcReadZeroCopyEnabled() != DEFAULT_OPTIONS.isGrpcReadZeroCopyEnabled()) { + gen.writeBooleanField("grpcReadZeroCopyEnabled", value.isGrpcReadZeroCopyEnabled()); + } + gen.writeEndObject(); + } +} + +class GcsReadOptionsDeserializer extends JsonDeserializer { + @Override + public GoogleCloudStorageReadOptions deserialize(JsonParser p, DeserializationContext ctxt) + throws java.io.IOException { + JsonNode root = p.readValueAsTree(); + GoogleCloudStorageReadOptions.Builder builder = + GcsReadOptionsSerializer.DEFAULT_OPTIONS.toBuilder(); + + if (root != null && root.isObject()) { + if (root.hasNonNull("fadvise")) { + builder.setFadvise( + GoogleCloudStorageReadOptions.Fadvise.valueOf(root.get("fadvise").asText())); + } + if (root.hasNonNull("fastFailOnNotFoundEnabled")) { + builder.setFastFailOnNotFoundEnabled(root.get("fastFailOnNotFoundEnabled").asBoolean()); + } else if (root.hasNonNull("fastFailOnNotFound")) { + builder.setFastFailOnNotFoundEnabled(root.get("fastFailOnNotFound").asBoolean()); + } + if (root.hasNonNull("minRangeRequestSize")) { + builder.setMinRangeRequestSize(root.get("minRangeRequestSize").asLong()); + } + if (root.hasNonNull("inplaceSeekLimit")) { + builder.setInplaceSeekLimit(root.get("inplaceSeekLimit").asLong()); + } + if (root.hasNonNull("grpcReadZeroCopyEnabled")) { + builder.setGrpcReadZeroCopyEnabled(root.get("grpcReadZeroCopyEnabled").asBoolean()); + } + } + return builder.build(); + } +} diff --git a/sdks/java/extensions/google-cloud-platform-core/src/test/java/org/apache/beam/sdk/extensions/gcp/GcpCoreApiSurfaceTest.java b/sdks/java/extensions/google-cloud-platform-core/src/test/java/org/apache/beam/sdk/extensions/gcp/GcpCoreApiSurfaceTest.java index 8af5e2260fc5..cd51d4fd9d28 100644 --- a/sdks/java/extensions/google-cloud-platform-core/src/test/java/org/apache/beam/sdk/extensions/gcp/GcpCoreApiSurfaceTest.java +++ b/sdks/java/extensions/google-cloud-platform-core/src/test/java/org/apache/beam/sdk/extensions/gcp/GcpCoreApiSurfaceTest.java @@ -51,6 +51,8 @@ public void testGcpCoreApiSurface() throws Exception { final Set>> allowedClasses = ImmutableSet.of( classesInPackage("com.fasterxml.jackson.annotation"), + classesInPackage("com.fasterxml.jackson.core"), + classesInPackage("com.fasterxml.jackson.databind"), classesInPackage("com.google.api.client.googleapis"), classesInPackage("com.google.api.client.http"), classesInPackage("com.google.api.client.json"), diff --git a/sdks/java/extensions/google-cloud-platform-core/src/test/java/org/apache/beam/sdk/extensions/gcp/options/GcsOptionsTest.java b/sdks/java/extensions/google-cloud-platform-core/src/test/java/org/apache/beam/sdk/extensions/gcp/options/GcsOptionsTest.java index c499290b851d..912f0110e9ce 100644 --- a/sdks/java/extensions/google-cloud-platform-core/src/test/java/org/apache/beam/sdk/extensions/gcp/options/GcsOptionsTest.java +++ b/sdks/java/extensions/google-cloud-platform-core/src/test/java/org/apache/beam/sdk/extensions/gcp/options/GcsOptionsTest.java @@ -18,11 +18,16 @@ package org.apache.beam.sdk.extensions.gcp.options; import static org.junit.Assert.assertEquals; +import static org.junit.Assert.assertFalse; +import static org.junit.Assert.assertNotNull; import static org.junit.Assert.assertThrows; +import com.fasterxml.jackson.databind.ObjectMapper; +import com.google.cloud.hadoop.gcsio.GoogleCloudStorageReadOptions; import java.util.Collections; import java.util.HashMap; import java.util.Map; +import org.apache.beam.sdk.options.PipelineOptions; import org.apache.beam.sdk.options.PipelineOptionsFactory; import org.junit.Test; import org.junit.runner.RunWith; @@ -75,4 +80,46 @@ public void testEntriesWithErrors() throws Exception { IllegalArgumentException.class, () -> PipelineOptionsFactory.fromArgs(TOO_MANY_ENTRIES_WITH_JOB).as(GcsOptions.class)); } + + @Test + public void testGoogleCloudStorageReadOptionsSerialization() throws Exception { + GcsOptions options = PipelineOptionsFactory.as(GcsOptions.class); + GoogleCloudStorageReadOptions readOptions = + GoogleCloudStorageReadOptions.builder() + .setFadvise(GoogleCloudStorageReadOptions.Fadvise.RANDOM) + .setFastFailOnNotFoundEnabled(false) + .setMinRangeRequestSize(12345L) + .build(); + options.setGoogleCloudStorageReadOptions(readOptions); + + ObjectMapper mapper = new ObjectMapper(); + String serialized = mapper.writeValueAsString(options); + GcsOptions deserialized = + mapper.readValue(serialized, PipelineOptions.class).as(GcsOptions.class); + + GoogleCloudStorageReadOptions deserializedReadOptions = + deserialized.getGoogleCloudStorageReadOptions(); + + assertNotNull(deserializedReadOptions); + assertEquals( + GoogleCloudStorageReadOptions.Fadvise.RANDOM, deserializedReadOptions.getFadvise()); + assertFalse(deserializedReadOptions.isFastFailOnNotFoundEnabled()); + assertEquals(12345L, deserializedReadOptions.getMinRangeRequestSize()); + } + + @Test + public void testDefaultGoogleCloudStorageReadOptionsSerialization() throws Exception { + GcsOptions options = PipelineOptionsFactory.as(GcsOptions.class); + ObjectMapper mapper = new ObjectMapper(); + String serialized = mapper.writeValueAsString(options); + GcsOptions deserialized = + mapper.readValue(serialized, PipelineOptions.class).as(GcsOptions.class); + + GoogleCloudStorageReadOptions deserializedReadOptions = + deserialized.getGoogleCloudStorageReadOptions(); + + assertNotNull(deserializedReadOptions); + assertEquals( + GoogleCloudStorageReadOptions.Fadvise.SEQUENTIAL, deserializedReadOptions.getFadvise()); + } } diff --git a/sdks/python/apache_beam/io/gcp/bigtableio_it_test.py b/sdks/python/apache_beam/io/gcp/bigtableio_it_test.py index 27b910ad5f08..488914c4b198 100644 --- a/sdks/python/apache_beam/io/gcp/bigtableio_it_test.py +++ b/sdks/python/apache_beam/io/gcp/bigtableio_it_test.py @@ -38,7 +38,7 @@ # Protect against environments where bigtable library is not available. try: - from apitools.base.py.exceptions import HttpError + from google.api_core.exceptions import GoogleAPICallError from google.cloud.bigtable import client from google.cloud.bigtable.row import Cell from google.cloud.bigtable.row import DirectRow @@ -48,7 +48,7 @@ from google.cloud.bigtable_admin_v2.types import instance except ImportError as e: client = None - HttpError = None + GoogleAPICallError = None def instance_prefix(instance): @@ -109,7 +109,7 @@ def tearDown(self): self.instance.instance_id) self.table.delete() self.instance.delete() - except HttpError: + except GoogleAPICallError: _LOGGER.warning( "Failed to clean up table [%s] and instance [%s]", self.table.table_id, @@ -208,7 +208,7 @@ def tearDown(self): try: _LOGGER.info("Deleting table [%s]", self.table.table_id) self.table.delete() - except HttpError: + except GoogleAPICallError: _LOGGER.warning("Failed to clean up table [%s]", self.table.table_id) @classmethod @@ -216,7 +216,7 @@ def tearDownClass(cls): try: _LOGGER.info("Deleting instance [%s]", cls.instance.instance_id) cls.instance.delete() - except HttpError: + except GoogleAPICallError: _LOGGER.warning( "Failed to clean up instance [%s]", cls.instance.instance_id) diff --git a/sdks/python/apache_beam/io/watch.py b/sdks/python/apache_beam/io/watch.py new file mode 100644 index 000000000000..2fcee7a8080f --- /dev/null +++ b/sdks/python/apache_beam/io/watch.py @@ -0,0 +1,730 @@ +# +# Licensed to the Apache Software Foundation (ASF) under one or more +# contributor license agreements. See the NOTICE file distributed with +# this work for additional information regarding copyright ownership. +# The ASF licenses this file to You under the Apache License, Version 2.0 +# (the "License"); you may not use this file except in compliance with +# the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# + +"""Experimental ``Watch`` transform for the Python SDK. + +``Watch`` continuously watches a growing set of outputs for each input element, +calling a user poll function on an interval until a per-input termination +condition fires. It is the engine behind periodic file-discovery and any +periodic polling source. + +For every input element the transform runs an independent loop:: + + poll -> keep never-seen-before outputs -> emit them (timestamped) -> + update watermark -> check termination -> wait(poll_interval) -> poll -> ... + +The output is an unbounded ``PCollection`` of ``(input, output)`` pairs. Each +output carries the event time the poll function first reported it. Dedup +hashes each output's key: the output itself by default, or +``output_key_fn(output)`` when one is given. The key coder is inferred when +not passed explicitly and converted to its deterministic form, so equal keys +hash equally across workers and restarts. + +Example:: + + from apache_beam.io.watch import Watch, PollResult, after_total_of + from apache_beam.transforms.window import TimestampedValue + from apache_beam.utils.timestamp import Duration, Timestamp + + def poll(prefix) -> PollResult[str]: + now = Timestamp.now() + outputs = [TimestampedValue(prefix + str(i), now) for i in range(3)] + return PollResult.complete(outputs) + + watched = inputs | Watch( + poll, + poll_interval=Duration(seconds=5), + termination=after_total_of(60)) + +This API is experimental and may change in backwards-incompatible ways. +""" + +import collections +import dataclasses +import hashlib +import inspect +import time +import typing +from typing import Any +from typing import Callable +from typing import Generic +from typing import Iterable +from typing import Optional +from typing import Tuple +from typing import TypeVar + +from apache_beam import coders +from apache_beam.coders.coders import Coder +from apache_beam.coders.coders import NullableCoder +from apache_beam.coders.coders import TimestampCoder +from apache_beam.coders.coders import TupleCoder +from apache_beam.io import iobase +from apache_beam.io.watermark_estimators import ManualWatermarkEstimator +from apache_beam.runners import sdf_utils +from apache_beam.transforms import PTransform +from apache_beam.transforms import core +from apache_beam.transforms.window import TimestampedValue +from apache_beam.utils.timestamp import MAX_TIMESTAMP +from apache_beam.utils.timestamp import Duration +from apache_beam.utils.timestamp import Timestamp + +__all__ = [ + 'Watch', + 'PollResult', + 'PollFn', + 'TerminationCondition', + 'never', + 'after_total_of', +] + +_HASH_DIGEST_SIZE = 16 # 128-bit digest width. + +OutputT = TypeVar('OutputT') + +# ------------------------------------------------------------------------------ +# Public API. +# ------------------------------------------------------------------------------ + + +@dataclasses.dataclass(frozen=True) +class PollResult(Generic[OutputT]): + """Outputs produced by one poll, plus an optional explicit watermark. + + ``watermark`` of ``None`` lets the transform infer the watermark from the + earliest new output. A watermark of ``MAX_TIMESTAMP`` (set by + :meth:`complete`) marks the input finished, so polling stops. + + The ``OutputT`` type parameter can annotate a poll function's return type, + as in ``-> PollResult[str]``; the transform infers the output coder from it. + """ + outputs: Tuple[TimestampedValue, ...] + watermark: Optional[Timestamp] = None + + @property + def is_complete(self) -> bool: + return self.watermark == MAX_TIMESTAMP + + @staticmethod + def _normalize(outputs, timestamp) -> Tuple[TimestampedValue, ...]: + if timestamp is None: + default_ts = Timestamp.now() + else: + default_ts = Timestamp.of(timestamp) + normalized = [] + for output in outputs: + if isinstance(output, TimestampedValue): + normalized.append(output) + else: + normalized.append(TimestampedValue(output, default_ts)) + return tuple(normalized) + + @staticmethod + def incomplete(outputs: Iterable, timestamp=None) -> 'PollResult': + """Reports outputs and expects more; the transform infers the watermark. + + A raw (non-:class:`TimestampedValue`) output is stamped with ``timestamp`` + when given, else with the current processing time. + """ + return PollResult(PollResult._normalize(outputs, timestamp), watermark=None) + + @staticmethod + def complete(outputs: Iterable, timestamp=None) -> 'PollResult': + """Reports the final outputs for an input, after which polling stops. + + A raw (non-:class:`TimestampedValue`) output is stamped with ``timestamp`` + when given, else with the current processing time. + """ + return PollResult( + PollResult._normalize(outputs, timestamp), watermark=MAX_TIMESTAMP) + + def with_watermark(self, watermark) -> 'PollResult': + return dataclasses.replace(self, watermark=Timestamp.of(watermark)) + + +class PollFn(object): + """Optional base for a poll function ``input -> PollResult``. + + Any callable with that signature works; subclass only to attach an output + coder hint via :meth:`default_output_coder`:: + + from apache_beam import coders + + class ListFiles(PollFn): + def __call__(self, prefix): + return PollResult.incomplete(list_files(prefix)) + + def default_output_coder(self): + return coders.StrUtf8Coder() + + A plain function can instead annotate its return type as ``PollResult[V]`` + and have the output coder inferred from ``V``. + """ + def __call__(self, element: Any) -> PollResult: + raise NotImplementedError + + def default_output_coder(self) -> Optional[Coder]: + return None + + +class TerminationCondition(object): + """Per-input stop policy with immutable, encodable state. + + Hooks follow the lifecycle of one input's polling loop. ``state`` flows from + :meth:`for_new_input` through the per-round hooks and is serialized with + :meth:`state_coder`. + """ + def for_new_input(self, now: Timestamp, element: Any) -> Any: + raise NotImplementedError + + def on_seen_new_output(self, now: Timestamp, state: Any) -> Any: + return state + + def on_poll_complete(self, state: Any) -> Any: + return state + + def can_stop_polling(self, now: Timestamp, state: Any) -> bool: + raise NotImplementedError + + def state_coder(self) -> Coder: + raise NotImplementedError + + +class _Never(TerminationCondition): + """Polls until the poll function returns :meth:`PollResult.complete`.""" + def for_new_input(self, now, element): + return 0 + + def can_stop_polling(self, now, state): + return False + + def state_coder(self): + return coders.VarIntCoder() + + +class _AfterTotalOf(TerminationCondition): + """Stops once the wall-clock time since the input was first seen exceeds a + fixed duration.""" + def __init__(self, duration: Duration): + self._duration_micros = duration.micros + + def for_new_input(self, now, element): + return (now, self._duration_micros) + + def can_stop_polling(self, now, state): + start, duration_micros = state + return (now - start).micros > duration_micros + + def state_coder(self): + return TupleCoder([TimestampCoder(), coders.VarIntCoder()]) + + +def never() -> TerminationCondition: + """Polls until :meth:`PollResult.complete`.""" + return _Never() + + +def after_total_of(duration) -> TerminationCondition: + """Stops polling an input after ``duration`` (a :class:`Duration` or seconds) + has elapsed since it was first seen.""" + return _AfterTotalOf(_as_duration(duration)) + + +# ------------------------------------------------------------------------------ +# Restriction state. +# ------------------------------------------------------------------------------ + + +class _GrowthState: + """Base for the two restriction variants a Watch input can hold.""" + + +@dataclasses.dataclass(frozen=True) +class _PollingGrowthState(_GrowthState): + """Keep-polling state: emitted-output hashes, watermark, termination state. + + ``completed`` maps a 16-byte output-key hash to the event time it was first + seen. It is insertion-ordered and treated as immutable; a new mapping is + built for each residual. + """ + completed: 'collections.OrderedDict[bytes, Timestamp]' + poll_watermark: Optional[Timestamp] + termination_state: Any + + +@dataclasses.dataclass(frozen=True) +class _NonPollingGrowthState(_GrowthState): + """Replay-then-stop state: the outputs already emitted this round. + + Produced as the checkpoint primary so a bundle retry re-emits exactly those + outputs. + """ + pending: PollResult + + +# Primary used when a checkpoint arrives before any claim; replays nothing. +_EMPTY_STATE = _NonPollingGrowthState(PollResult((), None)) + +# ------------------------------------------------------------------------------ +# Coders. +# ------------------------------------------------------------------------------ + + +class _TimestampedValueCoder(Coder): + """Coder for :class:`TimestampedValue`. + + ``TimestampedValue`` is normally unwrapped into a ``WindowedValue`` on the + wire, so the SDK ships no standalone coder for it. Watch keeps it inside the + restriction state, so this encodes the ``(value, timestamp)`` pair with a + :class:`TupleCoder` and rebuilds the ``TimestampedValue`` on decode. + """ + def __init__(self, value_coder: Coder): + self._tuple_coder = TupleCoder([value_coder, TimestampCoder()]) + + def encode(self, value: TimestampedValue) -> bytes: + return self._tuple_coder.encode((value.value, value.timestamp)) + + def decode(self, encoded: bytes) -> TimestampedValue: + value, timestamp = self._tuple_coder.decode(encoded) + return TimestampedValue(value, timestamp) + + def is_deterministic(self) -> bool: + return self._tuple_coder.is_deterministic() + + +class _GrowthStateCoder(Coder): + """Encodes a :class:`_PollingGrowthState` or :class:`_NonPollingGrowthState`. + + A ``(tag, payload)`` envelope selects the variant; the payload is a + variant-specific :class:`TupleCoder`. ``completed`` is encoded as an ordered + list of ``(hash, timestamp)`` pairs so insertion order survives a round trip. + This format is internal to the Python SDK. + """ + def __init__(self, output_coder: Coder, termination: TerminationCondition): + nullable_ts = NullableCoder(TimestampCoder()) + self._envelope_coder = TupleCoder( + [coders.VarIntCoder(), coders.BytesCoder()]) + self._polling_coder = TupleCoder([ + termination.state_coder(), + nullable_ts, + coders.ListCoder(TupleCoder([coders.BytesCoder(), TimestampCoder()])), + ]) + self._non_polling_coder = TupleCoder([ + nullable_ts, + coders.ListCoder(_TimestampedValueCoder(output_coder)), + ]) + + def encode(self, state: _GrowthState) -> bytes: + if isinstance(state, _PollingGrowthState): + payload = self._polling_coder.encode(( + state.termination_state, + state.poll_watermark, + list(state.completed.items()))) + return self._envelope_coder.encode((0, payload)) + payload = self._non_polling_coder.encode( + (state.pending.watermark, list(state.pending.outputs))) + return self._envelope_coder.encode((1, payload)) + + def decode(self, encoded: bytes) -> _GrowthState: + tag, payload = self._envelope_coder.decode(encoded) + if tag == 0: + termination_state, poll_watermark, items = self._polling_coder.decode( + payload) + return _PollingGrowthState( + collections.OrderedDict(items), poll_watermark, termination_state) + if tag == 1: + watermark, outputs = self._non_polling_coder.decode(payload) + return _NonPollingGrowthState(PollResult(tuple(outputs), watermark)) + raise ValueError('unknown Watch growth state tag: %r' % (tag, )) + + def is_deterministic(self) -> bool: + return False + + +# ------------------------------------------------------------------------------ +# Restriction tracker. +# ------------------------------------------------------------------------------ + + +def _identity(value: Any) -> Any: + return value + + +def _hash_output(key_coder: Coder, value: Any) -> bytes: + return hashlib.blake2b( + key_coder.encode(value), digest_size=_HASH_DIGEST_SIZE).digest() + + +def _max_watermark(left: Optional[Timestamp], + right: Optional[Timestamp]) -> Optional[Timestamp]: + if left is None: + return right + if right is None: + return left + return max(left, right) + + +def _never_seen_before( + restriction: _PollingGrowthState, + result: PollResult, + key_fn: Callable[[Any], Any], + key_coder: Coder) -> PollResult: + """Filters a poll result down to outputs whose key was never seen before. + + Dedup hashes ``key_fn(output.value)`` against the restriction's completed + set, also dropping in-round duplicates. Outputs are sorted by timestamp so + the earliest one can serve as the inferred watermark. + """ + new_outputs = [] + seen_this_round = set() + for output in result.outputs: + key_hash = _hash_output(key_coder, key_fn(output.value)) + if key_hash in restriction.completed or key_hash in seen_this_round: + continue + seen_this_round.add(key_hash) + new_outputs.append(output) + new_outputs.sort(key=lambda output: output.timestamp) + return dataclasses.replace(result, outputs=tuple(new_outputs)) + + +class _GrowthRestrictionTracker(iobase.RestrictionTracker): + """Tracks one input's polling restriction over claimed poll rounds. + + The claimed position is one poll round: a ``(PollResult, termination_state)`` + pair whose ``PollResult`` holds only never-seen-before outputs. ``process()`` + polls and dedups before claiming, so a slow poll never holds the tracker + lock; the tracker validates each claim against the restriction and derives + the checkpoint split from the claimed round in :meth:`try_split`. + """ + def __init__( + self, + restriction: _GrowthState, + key_fn: Callable[[Any], Any], + key_coder: Coder): + self._restriction = restriction + self._key_fn = key_fn + self._key_coder = key_coder + self._claimed_result = None # type: Optional[PollResult] + self._claimed_termination_state = None # type: Any + self._claimed_hashes = None # type: Optional[collections.OrderedDict] + self._should_stop = False + + def _hash(self, value: Any) -> bytes: + return _hash_output(self._key_coder, self._key_fn(value)) + + def current_restriction(self) -> _GrowthState: + return self._restriction + + def try_claim(self, position: Tuple[PollResult, Any]) -> bool: + """Claims one poll round; at most one claim succeeds per ``process()``. + + The claim is rejected after a checkpoint already stopped this invocation, + when a claimed output key was already completed, or when a replay does not + match the pending outputs exactly. + """ + if self._should_stop: + return False + result, termination_state = position + claimed_hashes = collections.OrderedDict() + for output in result.outputs: + claimed_hashes[self._hash(output.value)] = output.timestamp + if isinstance(self._restriction, _PollingGrowthState): + if any(key_hash in self._restriction.completed + for key_hash in claimed_hashes): + return False + else: + expected = set( + self._hash(output.value) + for output in self._restriction.pending.outputs) + if expected != set(claimed_hashes): + return False + self._should_stop = True + self._claimed_result = result + self._claimed_termination_state = termination_state + self._claimed_hashes = claimed_hashes + return True + + def try_split(self, fraction_of_remainder): + # Every split checkpoints at the claimed poll round; splitting a round + # further is not supported. + if self._claimed_result is None: + # No claim happened this invocation: the residual is all the work and + # the primary replays nothing. + residual = self._restriction + self._restriction = _EMPTY_STATE + elif isinstance(self._restriction, _NonPollingGrowthState): + # The claimed replay was the entire restriction, so nothing remains. + residual = _EMPTY_STATE + else: + # The primary becomes a replay of the claimed round; the residual + # resumes polling with the claimed keys marked completed. + merged = collections.OrderedDict(self._restriction.completed) + merged.update(self._claimed_hashes) + residual = _PollingGrowthState( + merged, + _max_watermark( + self._restriction.poll_watermark, self._claimed_result.watermark), + self._claimed_termination_state) + self._restriction = _NonPollingGrowthState(self._claimed_result) + self._should_stop = True + return self._restriction, residual + + def check_done(self) -> bool: + # Called after every process(); the single claim or a split sets the flag. + if self._should_stop: + return True + raise ValueError( + 'Watch restriction was neither claimed nor checkpointed: %r' % + (self._restriction, )) + + def current_progress(self) -> 'iobase.RestrictionProgress': + if self._should_stop: + return iobase.RestrictionProgress(completed=1.0, remaining=0.0) + return iobase.RestrictionProgress(completed=0.0, remaining=1.0) + + def is_bounded(self) -> bool: + # A polling restriction is unbounded; a replay-then-stop one is bounded. + return isinstance(self._restriction, _NonPollingGrowthState) + + +# ------------------------------------------------------------------------------ +# Splittable DoFn (its own restriction provider). +# ------------------------------------------------------------------------------ + + +class _WatchGrowthDoFn(core.DoFn, core.RestrictionProvider): + """Polling SDF that emits ``(input, output)`` pairs. + + The DoFn is its own ``RestrictionProvider``: ``RestrictionParam()`` with no + argument resolves the provider to the DoFn instance, so the provider methods + read the transform-level spec (poll function, coders, termination) off + ``self``. Provider methods run on a separately deserialized copy and before + ``setup()``, so the spec is immutable state set in ``__init__``. + """ + def __init__( + self, + poll_fn: Callable[[Any], PollResult], + termination: TerminationCondition, + poll_interval: Duration, + output_coder: Coder, + key_fn: Callable[[Any], Any], + key_coder: Coder, + now_fn: Optional[Callable[[], float]] = None): + self._poll_fn = poll_fn + self._termination = termination + self._poll_interval = poll_interval + self._output_coder = output_coder + self._key_fn = key_fn + self._key_coder = key_coder + self._now = now_fn or time.time + self._restriction_coder = _GrowthStateCoder(output_coder, termination) + + def initial_restriction(self, element) -> _PollingGrowthState: + now = Timestamp.of(self._now()) + return _PollingGrowthState( + collections.OrderedDict(), + None, + self._termination.for_new_input(now, element)) + + def create_tracker(self, restriction) -> _GrowthRestrictionTracker: + return _GrowthRestrictionTracker(restriction, self._key_fn, self._key_coder) + + def restriction_coder(self) -> Coder: + return self._restriction_coder + + def restriction_size(self, element, restriction) -> int: + return 1 + + @core.DoFn.unbounded_per_element() + def process( + self, + element, + timestamp=core.DoFn.TimestampParam, + tracker=core.DoFn.RestrictionParam(), + watermark_estimator=core.DoFn.WatermarkEstimatorParam( + ManualWatermarkEstimator.default_provider())): + assert isinstance(tracker, sdf_utils.RestrictionTrackerView) + # Java seeds the manual estimator with the element timestamp; the Python + # default provider starts at None, which a runner reads as MIN_TIMESTAMP + # and would pin the stage's output watermark until the first output. + if watermark_estimator.current_watermark() is None: + watermark_estimator.set_watermark(timestamp) + restriction = tracker.current_restriction() + if isinstance(restriction, _NonPollingGrowthState): + # Replay the outputs already emitted this round, then stop. No poll. + if not tracker.try_claim((restriction.pending, None)): + return + for output in restriction.pending.outputs: + yield TimestampedValue((element, output.value), output.timestamp) + return + # Poll before claiming so a slow poll never holds the tracker lock, which + # would block runner progress checks and checkpoints. + result = self._poll_fn(element) + # Read the clock after the poll so a slow poll counts against termination. + now = Timestamp.of(self._now()) + new_results = _never_seen_before( + restriction, result, self._key_fn, self._key_coder) + termination_state = restriction.termination_state + if new_results.outputs: + termination_state = self._termination.on_seen_new_output( + now, termination_state) + termination_state = self._termination.on_poll_complete(termination_state) + if not tracker.try_claim((new_results, termination_state)): + # A checkpoint already stopped this invocation; emit nothing. + return + for output in new_results.outputs: + yield TimestampedValue((element, output.value), output.timestamp) + if new_results.watermark is not None: + watermark = new_results.watermark + elif new_results.outputs: + # Outputs are timestamp-sorted, so the first one is the earliest. + watermark = new_results.outputs[0].timestamp + else: + watermark = None + if self._termination.can_stop_polling(now, termination_state): + return + if watermark is not None and watermark >= MAX_TIMESTAMP: + # No more output is possible (PollResult.complete), so polling stops. + return + if watermark is not None: + _set_watermark_if_greater(watermark_estimator, watermark) + tracker.defer_remainder(self._poll_interval) + + +def _set_watermark_if_greater(watermark_estimator, new_watermark) -> None: + # set_watermark raises on regression, so only ever advance the watermark. + current = watermark_estimator.current_watermark() + if current is None or new_watermark > current: + watermark_estimator.set_watermark(new_watermark) + + +# ------------------------------------------------------------------------------ +# Public PTransform. +# ------------------------------------------------------------------------------ + + +def _return_type(fn) -> Any: + """The return type annotation of ``fn`` or its ``__call__``, else ``Any``.""" + target = fn if inspect.isroutine(fn) else getattr(type(fn), '__call__', None) + if target is None: + return Any + try: + hints = typing.get_type_hints(target) + except (NameError, TypeError): + return Any + return hints.get('return', Any) + + +def _poll_output_type(poll_fn) -> Any: + """The ``V`` of a ``PollResult[V]`` return annotation on ``poll_fn``. + + This mirrors the Java SDK, which infers the output coder from the + ``PollFn``'s ``OutputT`` type parameter. Returns ``Any`` when ``poll_fn`` + carries no such annotation. + """ + hint = _return_type(poll_fn) + if typing.get_origin(hint) is PollResult: + args = typing.get_args(hint) + if len(args) == 1: + return args[0] + return Any + + +class Watch(PTransform): + """Watches a growing set of outputs per input via a periodic poll function. + + The output is an unbounded ``PCollection`` of ``(input, output)`` pairs. + + Args: + poll_fn: callable ``input -> PollResult``, invoked once per poll round. + poll_interval: delay between two poll rounds for one input, as a + :class:`Duration` or in seconds. + termination: per-input stop policy; defaults to :func:`never`. + output_coder: coder for the poll outputs, used to keep them in the + restriction state. Inferred when omitted: from a :class:`PollFn`'s + :meth:`~PollFn.default_output_coder`, else from the registered coder for + the ``V`` of a ``PollResult[V]`` return annotation on ``poll_fn``. + output_key_fn: derives the dedup key from an output; an output is emitted + only when its key was never seen before. Defaults to the output itself. + output_key_coder: coder whose encoding of the key is hashed for dedup; + inferred like ``output_coder`` when omitted. It is converted with + ``as_deterministic_coder`` so equal keys always hash equally; a coder + with no deterministic form is rejected. + now_fn: clock used for termination decisions; tests can inject one. + """ + def __init__( + self, + poll_fn: Callable[[Any], PollResult], + poll_interval, + termination: Optional[TerminationCondition] = None, + output_coder: Optional[Coder] = None, + output_key_fn: Optional[Callable[[Any], Any]] = None, + output_key_coder: Optional[Coder] = None, + now_fn: Optional[Callable[[], float]] = None): + super().__init__() + if poll_interval is None: + raise ValueError('Watch requires a poll_interval') + self._poll_fn = poll_fn + self._poll_interval = _as_duration(poll_interval) + self._termination = termination or never() + self._output_coder = output_coder + self._output_key_fn = output_key_fn + self._output_key_coder = output_key_coder + self._now = now_fn + + def expand(self, pcoll): + output_coder = self._output_coder + if output_coder is None and isinstance(self._poll_fn, PollFn): + output_coder = self._poll_fn.default_output_coder() + if output_coder is None: + output_coder = coders.registry.get_coder(_poll_output_type(self._poll_fn)) + if self._output_key_fn is None: + # The output is its own dedup key, so the key coder is the output coder. + key_fn = _identity + key_coder = self._output_key_coder or output_coder + else: + key_fn = self._output_key_fn + key_coder = self._output_key_coder or coders.registry.get_coder( + _return_type(self._output_key_fn)) + # Dedup hashes the encoded key, so equal keys must encode equally; use the + # coder's deterministic form and reject coders that have none. + key_coder = key_coder.as_deterministic_coder( + self.label, + 'Watch dedups by hashing the encoded output key, so the key coder ' + 'must be deterministic. %s has no deterministic form; pass a ' + 'deterministic output_key_coder (or output_coder).' % + type(key_coder).__name__) + # Type the (input, output) pairs from the input type and the resolved + # coder's type, so downstream transforms are typed and coder inference does + # not fall back to pickling. + input_type = pcoll.element_type or Any + try: + value_type = output_coder.to_type_hint() + except NotImplementedError: + value_type = Any + return pcoll | core.ParDo( + _WatchGrowthDoFn( + self._poll_fn, + self._termination, + self._poll_interval, + output_coder, + key_fn, + key_coder, + self._now)).with_output_types(Tuple[input_type, value_type]) + + +def _as_duration(value) -> Duration: + return value if isinstance(value, Duration) else Duration(value) diff --git a/sdks/python/apache_beam/io/watch_test.py b/sdks/python/apache_beam/io/watch_test.py new file mode 100644 index 000000000000..8c1f6571da66 --- /dev/null +++ b/sdks/python/apache_beam/io/watch_test.py @@ -0,0 +1,456 @@ +# +# Licensed to the Apache Software Foundation (ASF) under one or more +# contributor license agreements. See the NOTICE file distributed with +# this work for additional information regarding copyright ownership. +# The ASF licenses this file to You under the Apache License, Version 2.0 +# (the "License"); you may not use this file except in compliance with +# the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# + +"""Tests for the Watch transform.""" + +import collections +import unittest + +import apache_beam as beam +from apache_beam.coders.coders import Coder +from apache_beam.coders.coders import StrUtf8Coder +from apache_beam.io.watch import PollFn +from apache_beam.io.watch import PollResult +from apache_beam.io.watch import Watch +from apache_beam.io.watch import _GrowthRestrictionTracker +from apache_beam.io.watch import _GrowthStateCoder +from apache_beam.io.watch import _never_seen_before +from apache_beam.io.watch import _NonPollingGrowthState +from apache_beam.io.watch import _PollingGrowthState +from apache_beam.io.watch import _WatchGrowthDoFn +from apache_beam.io.watch import after_total_of +from apache_beam.io.watch import never +from apache_beam.io.watermark_estimators import ManualWatermarkEstimator +from apache_beam.options.pipeline_options import PipelineOptions +from apache_beam.runners.sdf_utils import RestrictionTrackerView +from apache_beam.runners.sdf_utils import ThreadsafeRestrictionTracker +from apache_beam.runners.sdf_utils import ThreadsafeWatermarkEstimator +from apache_beam.testing.test_pipeline import TestPipeline +from apache_beam.testing.util import TestWindowedValue +from apache_beam.testing.util import assert_that +from apache_beam.testing.util import equal_to +from apache_beam.transforms.window import FixedWindows +from apache_beam.transforms.window import GlobalWindow +from apache_beam.transforms.window import TimestampedValue +from apache_beam.typehints import typehints +from apache_beam.utils.timestamp import MAX_TIMESTAMP +from apache_beam.utils.timestamp import Duration +from apache_beam.utils.timestamp import Timestamp + + +def _ts(value, timestamp): + return TimestampedValue(value, Timestamp(timestamp)) + + +def _identity(output): + return output + + +def _new_results(restriction, result, key_fn=None): + return _never_seen_before( + restriction, result, key_fn or _identity, StrUtf8Coder()) + + +def _tracker(restriction): + return _GrowthRestrictionTracker(restriction, _identity, StrUtf8Coder()) + + +def _initial_polling(termination=None, now=Timestamp(0)): + termination = termination or never() + return _PollingGrowthState( + collections.OrderedDict(), None, termination.for_new_input(now, 'input')) + + +class GrowthStateCoderTest(unittest.TestCase): + def test_polling_round_trip_preserves_resume_state(self): + termination = after_total_of(Duration(30)) + coder = _GrowthStateCoder(StrUtf8Coder(), termination) + completed = collections.OrderedDict([ + (b'a' * 16, Timestamp(1)), + (b'b' * 16, Timestamp(2)), + (b'c' * 16, Timestamp(3)), + ]) + termination_state = termination.for_new_input(Timestamp(7), 'input') + state = _PollingGrowthState(completed, Timestamp(5), termination_state) + decoded = coder.decode(coder.encode(state)) + self.assertEqual(list(completed.items()), list(decoded.completed.items())) + self.assertEqual(Timestamp(5), decoded.poll_watermark) + self.assertEqual(termination_state, decoded.termination_state) + + def test_non_polling_round_trip_preserves_pending_outputs(self): + coder = _GrowthStateCoder(StrUtf8Coder(), never()) + pending = PollResult((_ts('a', 1), _ts('b', 2)), MAX_TIMESTAMP) + state = _NonPollingGrowthState(pending) + decoded = coder.decode(coder.encode(state)) + self.assertEqual(MAX_TIMESTAMP, decoded.pending.watermark) + self.assertEqual([('a', Timestamp(1)), ('b', Timestamp(2))], + [(o.value, o.timestamp) for o in decoded.pending.outputs]) + + +class NeverSeenBeforeTest(unittest.TestCase): + def test_dedups_and_sorts_by_timestamp(self): + result = PollResult.incomplete([_ts('b', 2), _ts('a', 1), _ts('a', 1)]) + new_results = _new_results(_initial_polling(), result) + self.assertEqual(['a', 'b'], [o.value for o in new_results.outputs]) + + def test_dedups_against_completed_keys(self): + state = _initial_polling() + first = _new_results( + state, PollResult.incomplete([_ts('a', 1), _ts('b', 2)])) + tracker = _tracker(state) + self.assertTrue(tracker.try_claim((first, 0))) + _, residual = tracker.try_split(0) + second = _new_results( + residual, PollResult.incomplete([_ts('a', 1), _ts('c', 3)])) + self.assertEqual(['c'], [o.value for o in second.outputs]) + + def test_output_key_dedups_by_derived_key(self): + result = PollResult.incomplete([_ts('a1', 1), _ts('a2', 2), _ts('b1', 3)]) + # The key is the first character, so 'a1' and 'a2' collapse to one output. + new_results = _new_results( + _initial_polling(), result, key_fn=lambda output: output[0]) + self.assertEqual(['a1', 'b1'], [o.value for o in new_results.outputs]) + + def test_preserves_explicit_watermark(self): + result = PollResult.incomplete([_ts('c', 3)]).with_watermark(5) + new_results = _new_results(_initial_polling(), result) + self.assertEqual(Timestamp(5), new_results.watermark) + + +class GrowthTrackerTest(unittest.TestCase): + def test_claim_then_split_builds_replay_primary_and_merged_residual(self): + state = _initial_polling() + new_results = _new_results( + state, PollResult.incomplete([_ts('a', 1), _ts('b', 2)])) + tracker = _tracker(state) + self.assertFalse(tracker.is_bounded()) + self.assertTrue(tracker.try_claim((new_results, 0))) + primary, residual = tracker.try_split(0) + self.assertIsInstance(primary, _NonPollingGrowthState) + self.assertEqual(new_results, primary.pending) + self.assertIsInstance(residual, _PollingGrowthState) + self.assertEqual(2, len(residual.completed)) + self.assertEqual(0, residual.termination_state) + self.assertTrue(tracker.check_done()) + + def test_split_merges_explicit_watermark_into_residual(self): + state = _initial_polling() + result = PollResult.incomplete([_ts('c', 3)]).with_watermark(5) + tracker = _tracker(state) + self.assertTrue(tracker.try_claim((_new_results(state, result), 0))) + _, residual = tracker.try_split(0) + self.assertEqual(Timestamp(5), residual.poll_watermark) + + def test_second_claim_is_rejected(self): + state = _initial_polling() + new_results = _new_results(state, PollResult.incomplete([_ts('a', 1)])) + tracker = _tracker(state) + self.assertTrue(tracker.try_claim((new_results, 0))) + self.assertFalse(tracker.try_claim((new_results, 0))) + + def test_claim_rejects_already_completed_keys(self): + # The tracker re-validates a claim, so a poll round that was not deduped + # against the restriction is rejected instead of emitting duplicates. + state = _initial_polling() + first = _new_results(state, PollResult.incomplete([_ts('a', 1)])) + tracker = _tracker(state) + self.assertTrue(tracker.try_claim((first, 0))) + _, residual = tracker.try_split(0) + stale = PollResult.incomplete([_ts('a', 1)]) + self.assertFalse(_tracker(residual).try_claim((stale, 0))) + + def test_split_before_claim_moves_all_work_to_residual(self): + state = _initial_polling() + tracker = _tracker(state) + primary, residual = tracker.try_split(0) + self.assertIs(state, residual) + self.assertIsInstance(primary, _NonPollingGrowthState) + self.assertEqual((), primary.pending.outputs) + new_results = _new_results(state, PollResult.incomplete([_ts('a', 1)])) + self.assertFalse(tracker.try_claim((new_results, 0))) + self.assertTrue(tracker.check_done()) + + def test_non_polling_replays_exactly_the_pending_outputs(self): + pending = PollResult((_ts('a', 1), _ts('b', 2)), MAX_TIMESTAMP) + tracker = _tracker(_NonPollingGrowthState(pending)) + self.assertTrue(tracker.is_bounded()) + # A replay must claim the pending poll result exactly. + partial = PollResult((_ts('a', 1), ), None) + self.assertFalse(tracker.try_claim((partial, None))) + self.assertTrue(tracker.try_claim((pending, None))) + # A checkpoint after the replay leaves no residual work. + _, residual = tracker.try_split(0) + self.assertEqual((), residual.pending.outputs) + self.assertTrue(tracker.check_done()) + + def test_check_done_raises_without_claim_or_split(self): + tracker = _tracker(_initial_polling()) + with self.assertRaises(ValueError): + tracker.check_done() + + def test_wrapper_chain_defers_merged_residual(self): + state = _initial_polling() + new_results = _new_results( + state, PollResult.incomplete([_ts('a', 1), _ts('b', 2)])) + threadsafe = ThreadsafeRestrictionTracker(_tracker(state)) + view = RestrictionTrackerView(threadsafe) + self.assertTrue(view.try_claim((new_results, 0))) + view.defer_remainder(Duration(5)) + residual, _ = threadsafe.deferred_status() + self.assertIsInstance(residual, _PollingGrowthState) + self.assertEqual(2, len(residual.completed)) + + +class TerminationConditionTest(unittest.TestCase): + def test_never_does_not_stop(self): + termination = never() + state = termination.for_new_input(Timestamp(0), 'input') + self.assertFalse(termination.can_stop_polling(MAX_TIMESTAMP, state)) + + def test_after_total_of_stops_once_duration_elapsed(self): + termination = after_total_of(10) + state = termination.for_new_input(Timestamp(0), 'input') + self.assertFalse(termination.can_stop_polling(Timestamp(10), state)) + self.assertTrue(termination.can_stop_polling(Timestamp(11), state)) + + +# Module-level so the poll function pickles by reference; the call counter is +# shared within the single in-memory DirectRunner process. +_POLL_CALLS = collections.defaultdict(int) + + +def _growing_poll(prefix): + # Unannotated on purpose: dedup must hold on the inferred fallback coder. + _POLL_CALLS[prefix] += 1 + count = _POLL_CALLS[prefix] + outputs = [_ts('%s%d' % (prefix, i), i + 1) for i in range(count)] + if count >= 3: + return PollResult.complete(outputs) + return PollResult.incomplete(outputs) + + +def _complete_poll(prefix) -> PollResult[str]: + return PollResult.complete([_ts(prefix + 'a', 1), _ts(prefix + 'b', 2)]) + + +def _first_char(output): + return output[0] + + +def _empty_poll(unused_element): + return PollResult.incomplete([]) + + +def _keyed_poll(prefix): + # 'a1' and 'a2' share the dedup key 'a', so only 'a1' is emitted. + return PollResult.complete([_ts('a1', 1), _ts('a2', 2), _ts('b1', 3)]) + + +class _StrCoderPollFn(PollFn): + def __call__(self, element): + return PollResult.complete([_ts(element + 'a', 1)]) + + def default_output_coder(self): + return StrUtf8Coder() + + +class _NoDeterministicFormCoder(Coder): + def encode(self, value): + return b'' + + def decode(self, encoded): + return None + + def is_deterministic(self): + return False + + +def _windowed_group(kv, window=beam.DoFn.WindowParam): + return ((window.start, window.end), sorted(kv[1])) + + +class WatchDoFnProcessTest(unittest.TestCase): + def _process( + self, poll_fn, element, timestamp, restriction=None, watermark=None): + dofn = _WatchGrowthDoFn( + poll_fn, + never(), + Duration(1), + StrUtf8Coder(), + _identity, + StrUtf8Coder()) + if restriction is None: + restriction = dofn.initial_restriction(element) + threadsafe = ThreadsafeRestrictionTracker(dofn.create_tracker(restriction)) + estimator = ThreadsafeWatermarkEstimator( + ManualWatermarkEstimator(watermark)) + outputs = list( + dofn.process( + element, + timestamp=timestamp, + tracker=RestrictionTrackerView(threadsafe), + watermark_estimator=estimator)) + return outputs, threadsafe, estimator + + def test_empty_round_holds_watermark_at_input_timestamp(self): + outputs, threadsafe, estimator = self._process( + _empty_poll, 'in', Timestamp(7)) + self.assertEqual([], outputs) + # The estimator is seeded from the input timestamp, so the deferred + # residual holds the watermark there instead of at MIN_TIMESTAMP. + self.assertEqual(Timestamp(7), estimator.current_watermark()) + residual, _ = threadsafe.deferred_status() + self.assertIsInstance(residual, _PollingGrowthState) + + def test_complete_round_stops_without_residual(self): + outputs, threadsafe, _ = self._process(_complete_poll, 'k:', Timestamp(0)) + self.assertEqual([('k:', 'k:a'), ('k:', 'k:b')], + [value.value for value in outputs]) + self.assertIsNone(threadsafe.deferred_status()) + self.assertTrue(threadsafe.check_done()) + + def test_replay_round_leaves_the_watermark_alone(self): + pending = PollResult((_ts('k:a', 1), _ts('k:b', 2)), MAX_TIMESTAMP) + outputs, threadsafe, estimator = self._process( + _empty_poll, + 'k:', + Timestamp(7), + restriction=_NonPollingGrowthState(pending)) + self.assertEqual([('k:', 'k:a'), ('k:', 'k:b')], + [value.value for value in outputs]) + # The replay branch holds the watermark at the seed, so it never runs ahead + # of the replayed outputs and never releases to MAX_TIMESTAMP itself. + self.assertEqual(Timestamp(7), estimator.current_watermark()) + self.assertIsNone(threadsafe.deferred_status()) + self.assertTrue(threadsafe.check_done()) + + def test_terminal_round_after_deferring_leaves_no_residual(self): + _POLL_CALLS.clear() + # Round one defers and parks the watermark on the new output's time. + _, threadsafe, estimator = self._process(_growing_poll, 'd:', Timestamp(0)) + residual, _ = threadsafe.deferred_status() + self.assertIsInstance(residual, _PollingGrowthState) + self.assertEqual(Timestamp(1), estimator.current_watermark()) + # Round two resumes from that residual, carrying the watermark forward. + _, threadsafe, estimator = self._process( + _growing_poll, + 'd:', + Timestamp(0), + restriction=residual, + watermark=estimator.current_watermark()) + residual, _ = threadsafe.deferred_status() + self.assertEqual(Timestamp(2), estimator.current_watermark()) + # Round three completes. The watermark stays where round two left it and + # the round reports no residual, so nothing carries that hold forward. + outputs, threadsafe, estimator = self._process( + _growing_poll, + 'd:', + Timestamp(0), + restriction=residual, + watermark=estimator.current_watermark()) + self.assertEqual([('d:', 'd:2')], [value.value for value in outputs]) + self.assertEqual(Timestamp(2), estimator.current_watermark()) + self.assertIsNone(threadsafe.deferred_status()) + self.assertTrue(threadsafe.check_done()) + + +class WatchEndToEndTest(unittest.TestCase): + def _in_memory_pipeline(self): + return TestPipeline( + options=PipelineOptions(direct_running_mode='in_memory')) + + def test_complete_outputs_values_and_timestamps(self): + with self._in_memory_pipeline() as p: + output = ( + p | beam.Create(['k:']) + | Watch(_complete_poll, poll_interval=Duration(1))) + assert_that( + output, + equal_to([ + TestWindowedValue(('k:', 'k:a'), Timestamp(1), [GlobalWindow()]), + TestWindowedValue(('k:', 'k:b'), Timestamp(2), [GlobalWindow()]), + ]), + reify_windows=True) + + def test_complete_advances_watermark_for_windowed_pipeline(self): + with self._in_memory_pipeline() as p: + output = ( + p | beam.Create(['k:']) + | Watch(_complete_poll, poll_interval=Duration(1))) + grouped = ( + output + | beam.WindowInto(FixedWindows(10)) + | beam.Map(lambda kv: ('all', kv[1])) + | beam.GroupByKey() + | beam.Map(_windowed_group)) + assert_that( + grouped, + equal_to([ + ((Timestamp(0), Timestamp(10)), ['k:a', 'k:b']), + ])) + + def test_multi_round_dedups_stops_and_is_per_input(self): + _POLL_CALLS.clear() + with self._in_memory_pipeline() as p: + output = ( + p | beam.Create(['x:', 'y:']) + | Watch(_growing_poll, poll_interval=Duration(0.05))) + assert_that( + output, + equal_to([('x:', 'x:0'), ('x:', 'x:1'), ('x:', 'x:2'), ('y:', 'y:0'), + ('y:', 'y:1'), ('y:', 'y:2')])) + self.assertEqual(3, _POLL_CALLS['x:']) + self.assertEqual(3, _POLL_CALLS['y:']) + + def test_output_key_dedups_across_pipeline(self): + with self._in_memory_pipeline() as p: + output = ( + p | beam.Create(['k']) + | Watch( + _keyed_poll, poll_interval=Duration(1), + output_key_fn=_first_char)) + assert_that(output, equal_to([('k', 'a1'), ('k', 'b1')])) + + def test_rejects_key_coder_without_deterministic_form(self): + with self.assertRaises(ValueError): + with self._in_memory_pipeline() as p: + _ = ( + p | beam.Create(['k:']) + | Watch( + _complete_poll, + poll_interval=Duration(1), + output_key_coder=_NoDeterministicFormCoder())) + + def test_infers_output_coder_from_return_annotation(self): + # _complete_poll is annotated ``-> PollResult[str]``, so the output coder + # and with it the (input, output) element type are inferred without hints. + with self._in_memory_pipeline() as p: + output = ( + p | beam.Create(['k:']) + | Watch(_complete_poll, poll_interval=Duration(1))) + self.assertEqual(typehints.Tuple[str, str], output.element_type) + + def test_uses_poll_fn_default_output_coder(self): + with self._in_memory_pipeline() as p: + output = ( + p | beam.Create(['k:']) + | Watch(_StrCoderPollFn(), poll_interval=Duration(1))) + self.assertEqual(typehints.Tuple[str, str], output.element_type) + + +if __name__ == '__main__': + unittest.main() diff --git a/sdks/python/apache_beam/testing/util.py b/sdks/python/apache_beam/testing/util.py index cfe8221b4aa2..bf969d5aa00f 100644 --- a/sdks/python/apache_beam/testing/util.py +++ b/sdks/python/apache_beam/testing/util.py @@ -22,6 +22,8 @@ import collections import glob import io +import math +import numbers import tempfile from typing import Any from typing import Iterable @@ -44,6 +46,7 @@ __all__ = [ 'assert_that', 'equal_to', + 'equal_to_approx', 'equal_to_per_window', 'has_at_least_one', 'is_empty', @@ -231,6 +234,46 @@ def row_namedtuple_equals_fn(expected, actual, fallback_equals_fn=None): return True +def equal_to_approx(expected, rel_tol=1e-09, abs_tol=0.0): + """Matcher used by assert_that that compares numeric elements approximately. + + Behaves similarly to `equal_to` for sequence ordering and membership, but any + real number elements (integers and floats) are compared using `math.isclose` + instead of exact equality. + + Approximate comparisons are also applied to real numbers nested within lists + and tuples. Note that `equal_to`'s advanced handling for Beam `Row` and + `NamedTuple` is not supported here. All other elements are compared using + standard `==` equality. + + Args: + expected: The expected output or sequence to compare against. + rel_tol: The relative tolerance used by `math.isclose` (default: 1e-09). + abs_tol: The absolute tolerance used by `math.isclose` (default: 0.0). + + Example: + assert_that( + pipeline | beam.Create([1.000000001, 2.0]), + equal_to_approx([1.0, 2.0])) + """ + def _approx_equals(expected_element, actual_element): + return _elements_approx_equal( + expected_element, actual_element, rel_tol, abs_tol) + + return equal_to(expected, equals_fn=_approx_equals) + + +def _elements_approx_equal(expected, actual, rel_tol, abs_tol): + if all(isinstance(x, numbers.Real) for x in (expected, actual)): + return math.isclose(expected, actual, rel_tol=rel_tol, abs_tol=abs_tol) + if (isinstance(expected, (list, tuple)) and type(actual) is type(expected) and + len(expected) == len(actual)): + return all( + _elements_approx_equal(e, a, rel_tol, abs_tol) + for e, a in zip(expected, actual)) + return expected == actual + + def matches_all(expected): """Matcher used by assert_that to check a set of matchers. diff --git a/sdks/python/apache_beam/testing/util_test.py b/sdks/python/apache_beam/testing/util_test.py index 12314f4653aa..87f6162ea311 100644 --- a/sdks/python/apache_beam/testing/util_test.py +++ b/sdks/python/apache_beam/testing/util_test.py @@ -29,6 +29,7 @@ from apache_beam.testing.util import TestWindowedValue from apache_beam.testing.util import assert_that from apache_beam.testing.util import equal_to +from apache_beam.testing.util import equal_to_approx from apache_beam.testing.util import equal_to_per_window from apache_beam.testing.util import is_empty from apache_beam.testing.util import is_not_empty @@ -93,6 +94,49 @@ def test_assert_with_custom_comparator(self): p | Create([1, 2, 3]), equal_to(['1', '2', '3'], equals_fn=lambda e, a: int(e) == int(a))) + def test_equal_to_approx(self): + with TestPipeline() as p: + assert_that( + p | Create([1.0, 2, 3.0]), equal_to_approx([3.0000000001, 2.0, 1.0])) + + def test_equal_to_approx_nested(self): + with TestPipeline() as p: + assert_that( + p | Create([('a', 1.0), ('b', 2.0)]), + equal_to_approx([('b', 2.0000000001), ('a', 1)])) + + def test_equal_to_approx_with_abs_tol(self): + with TestPipeline() as p: + assert_that(p | Create([0.0]), equal_to_approx([1e-10], abs_tol=1e-9)) + + def test_equal_to_approx_fails_outside_tolerance(self): + with self.assertRaises(Exception): + with TestPipeline() as p: + assert_that(p | Create([1.0]), equal_to_approx([1.1])) + + def test_equal_to_approx_nested_list(self): + with TestPipeline() as p: + assert_that( + p | Create([[1.0, 2.0]]), equal_to_approx([[1.0000000001, 2.0]])) + + def test_equal_to_approx_non_numeric(self): + with TestPipeline() as p: + assert_that(p | Create(['a', 'b']), equal_to_approx(['b', 'a'])) + + def test_equal_to_approx_empty(self): + with TestPipeline() as p: + assert_that(p | Create([]), equal_to_approx([])) + + def test_equal_to_approx_with_rel_tol(self): + with TestPipeline() as p: + assert_that( + p | Create([100.0]), equal_to_approx([100.00001], rel_tol=1e-6)) + + def test_equal_to_approx_nested_fails_outside_tolerance(self): + with self.assertRaises(Exception): + with TestPipeline() as p: + assert_that(p | Create([('a', 1.0)]), equal_to_approx([('a', 1.2)])) + def test_reified_value_passes(self): expected = [ TestWindowedValue(v, MIN_TIMESTAMP, [GlobalWindow()])