diff --git a/CHANGES.md b/CHANGES.md index 8fa59e7d5ea0..70218d411432 100644 --- a/CHANGES.md +++ b/CHANGES.md @@ -139,6 +139,7 @@ * (Python) Setting `max_batch_weight` on a RunInference `ModelHandler` no longer fails with a `TypeError`: `BatchElements` now supports `max_batch_weight`, which limits the total `element_size_fn` of a batch in addition to its number of elements ([#40644](https://github.com/apache/beam/issues/40644)). * (Python) `BatchElements` now honors `element_size_fn` when `max_batch_duration_secs` is set, and no longer emits an empty batch before an element that exceeds the batch size on its own ([#40644](https://github.com/apache/beam/issues/40644)). * (Python) `KeyedModelHandler` and `MaybeKeyedModelHandler` now apply the `element_size_fn` and `batch_length_fn` of the wrapped `ModelHandler` to examples, instead of to `(key, example)` pairs ([#40644](https://github.com/apache/beam/issues/40644)). +* (Java) Portable runners no longer deadlock when a single bundle produces more output than the Fn API data buffers hold (about 1 GB with the Python SDK, e.g. one Flink batch partition). The runner now consumes a bundle's output before waiting for its `ProcessBundleResponse` ([#40720](https://github.com/apache/beam/issues/40720)). * (Python) DataFrame API: `drop_duplicates`, `duplicated` and `nunique` no longer fail on rows with NaN values or on categorical columns with unused categories, and `nunique(dropna=False)` now counts NaN ([#40570](https://github.com/apache/beam/issues/40570)). * (Java) BigQuery Storage Write API `TableRow` conversion now honors `allowMissingRequiredFields` when building protocol-buffer messages, so callers that pass `true` with required descriptors get partial messages instead of a conversion-time `UninitializedMessageException` ([#40189](https://github.com/apache/beam/issues/40189)). * Fixed X (Java/Python) ([#X](https://github.com/apache/beam/issues/X)). diff --git a/runners/java-fn-execution/src/main/java/org/apache/beam/runners/fnexecution/control/SdkHarnessClient.java b/runners/java-fn-execution/src/main/java/org/apache/beam/runners/fnexecution/control/SdkHarnessClient.java index 704d298a195d..da3d98118751 100644 --- a/runners/java-fn-execution/src/main/java/org/apache/beam/runners/fnexecution/control/SdkHarnessClient.java +++ b/runners/java-fn-execution/src/main/java/org/apache/beam/runners/fnexecution/control/SdkHarnessClient.java @@ -516,22 +516,58 @@ public void close() throws Exception { } finally { beamFnDataOutboundAggregator.discard(); } + // Consume the bundle's output BEFORE waiting for the ProcessBundleResponse. The inbound + // queue is bounded, and the SDK does not finish the bundle until its output has been sent. + // Waiting for the response first therefore deadlocks once a single bundle produces more + // output than the inbound queue and the SDK's own send buffer hold together: the SDK + // blocks sending, the response never arrives, and nothing drains the queue. Batch runners + // that hand a whole partition to the SDK as one bundle hit this with ~1 GB of output. + boolean inboundDrained = false; + Exception drainException = null; + if (exception == null && beamFnDataInboundObserver.isPresent()) { + BeamFnDataInboundObserver inboundObserver = beamFnDataInboundObserver.get(); + // Async: when the response is already complete, a synchronous callback would run on + // this thread and could block putting elements into the queue this thread must drain. + response.whenCompleteAsync( + (completedResponse, failure) -> { + try { + if (failure != null) { + // The SDK will not send the rest of the output; unblock the drain below. + inboundObserver.close(); + } else if (completedResponse.hasElements()) { + inboundObserver.accept(completedResponse.getElements()); + } + } catch (Exception e) { + // Only possible once the drain has failed and cancelled the queue, and that + // failure is reported by close(). + LOG.debug("Dropping the response elements of failed bundle {}", bundleId, e); + } + }); + inboundDrained = true; + try { + inboundObserver.awaitCompletion(); + fnApiDataService.unregisterReceiver(bundleId); + } catch (Exception e) { + drainException = e; + } + } try { // We don't have to worry about the completion stage. if (exception == null) { BeamFnApi.ProcessBundleResponse completedResponse = MoreFutures.get(response); - if (completedResponse.hasElements()) { - beamFnDataInboundObserver.get().accept(completedResponse.getElements()); - } outstandingRequests.arriveAndAwaitAdvance(); - progressHandler.onCompleted(completedResponse); - if (completedResponse.getResidualRootsCount() > 0) { - checkpointHandler.onCheckpoint(completedResponse); - } - if (completedResponse.getRequiresFinalization()) { - finalizationHandler.requestsFinalization(bundleId); + // A bundle whose output could not be consumed has failed even if the SDK finished + // it, so it must not report completion or register residuals or finalization. + if (drainException == null) { + progressHandler.onCompleted(completedResponse); + if (completedResponse.getResidualRootsCount() > 0) { + checkpointHandler.onCheckpoint(completedResponse); + } + if (completedResponse.getRequiresFinalization()) { + finalizationHandler.requestsFinalization(bundleId); + } } } else { // TODO: [https://github.com/apache/beam/issues/18756] Handle aborting the bundle being @@ -546,6 +582,15 @@ public void close() throws Exception { exception.addSuppressed(e); } } + // After the response, so that a failed bundle reports the SDK's error rather than the + // cancelled queue it caused. + if (drainException != null) { + if (exception == null) { + exception = drainException; + } else { + exception.addSuppressed(drainException); + } + } try { if (exception == null) { stateRegistration.deregister(); @@ -559,14 +604,10 @@ public void close() throws Exception { exception.addSuppressed(e); } } - if (beamFnDataInboundObserver.isPresent()) { + if (beamFnDataInboundObserver.isPresent() && !inboundDrained) { + // Sending the bundle's input failed, so its output was never consumed. try { - if (exception == null) { - beamFnDataInboundObserver.get().awaitCompletion(); - fnApiDataService.unregisterReceiver(bundleId); - } else { - beamFnDataInboundObserver.get().close(); - } + beamFnDataInboundObserver.get().close(); } catch (Exception e) { if (exception == null) { exception = e; diff --git a/runners/java-fn-execution/src/test/java/org/apache/beam/runners/fnexecution/control/SdkHarnessClientTest.java b/runners/java-fn-execution/src/test/java/org/apache/beam/runners/fnexecution/control/SdkHarnessClientTest.java index 7b19376426cc..5dc6afe2808f 100644 --- a/runners/java-fn-execution/src/test/java/org/apache/beam/runners/fnexecution/control/SdkHarnessClientTest.java +++ b/runners/java-fn-execution/src/test/java/org/apache/beam/runners/fnexecution/control/SdkHarnessClientTest.java @@ -43,7 +43,9 @@ import java.util.List; import java.util.concurrent.CompletableFuture; import java.util.concurrent.ExecutionException; +import java.util.concurrent.ExecutorService; import java.util.concurrent.Executors; +import java.util.concurrent.Future; import java.util.concurrent.ScheduledExecutorService; import java.util.concurrent.TimeUnit; import org.apache.beam.model.fnexecution.v1.BeamFnApi; @@ -965,6 +967,186 @@ public void testBundleFinalizationCallback() throws Exception { verifyNoMoreInteractions(mockCheckpointHandler, mockSplitHandler); } + /** Output of {@code count} byte arrays for {@code bundleId}, one per data message. */ + private static List outputMessages(String bundleId, int count) throws Exception { + List messages = new ArrayList<>(); + for (int i = 0; i < count; i++) { + ByteString.Output data = ByteString.newOutput(); + ByteArrayCoder.of().encode(new byte[] {(byte) i}, data); + messages.add( + Elements.newBuilder() + .addData( + Elements.Data.newBuilder() + .setInstructionId(bundleId) + .setTransformId(SDK_GRPC_WRITE_TRANSFORM) + .setData(data.toByteString())) + .build()); + } + return messages; + } + + private static Elements lastOutputMessage(String bundleId) { + return Elements.newBuilder() + .addData( + Elements.Data.newBuilder() + .setInstructionId(bundleId) + .setTransformId(SDK_GRPC_WRITE_TRANSFORM) + .setIsLast(true)) + .build(); + } + + @Test + public void testOutputLargerThanTheInboundQueueDoesNotDeadlock() throws Exception { + // The SDK only completes a bundle once all of its output has been sent. If close() waits for + // the ProcessBundleResponse before consuming that output, a bundle whose output does not fit + // in the inbound queue (100 messages) never finishes. + CompletableFuture processBundleResponseFuture = new CompletableFuture<>(); + when(fnApiControlClient.handle(any(BeamFnApi.InstructionRequest.class))) + .thenReturn(processBundleResponseFuture); + FullWindowedValueCoder coder = + FullWindowedValueCoder.of(StringUtf8Coder.of(), Coder.INSTANCE); + BundleProcessor processor = + sdkHarnessClient.getProcessor( + descriptor, + Collections.singletonList( + RemoteInputDestination.of( + (FullWindowedValueCoder) coder, SDK_GRPC_READ_TRANSFORM))); + doNothing().when(dataService).registerReceiver(any(), outputReceiverCaptor.capture()); + when(dataService.createOutboundAggregator(any(), anyBoolean())) + .thenReturn(mock(BeamFnDataOutboundAggregator.class)); + + List received = Collections.synchronizedList(new ArrayList<>()); + ActiveBundle activeBundle = + processor.newBundle( + ImmutableMap.of( + SDK_GRPC_WRITE_TRANSFORM, + RemoteOutputReceiver.of(ByteArrayCoder.of(), received::add)), + BundleProgressHandler.ignored()); + String bundleId = activeBundle.getId(); + + ExecutorService executor = Executors.newFixedThreadPool(2); + try { + Future sdkSendsOutputThenResponse = + executor.submit( + () -> { + for (Elements message : outputMessages(bundleId, 150)) { + outputReceiverCaptor.getValue().accept(message); + } + outputReceiverCaptor.getValue().accept(lastOutputMessage(bundleId)); + processBundleResponseFuture.complete( + InstructionResponse.newBuilder() + .setProcessBundle(ProcessBundleResponse.getDefaultInstance()) + .build()); + return null; + }); + + Future closed = + executor.submit( + () -> { + activeBundle.close(); + return null; + }); + // Without the fix this times out: the SDK blocks on message 101 and close() waits forever. + closed.get(60, TimeUnit.SECONDS); + sdkSendsOutputThenResponse.get(60, TimeUnit.SECONDS); + } finally { + executor.shutdownNow(); + } + assertEquals(150, received.size()); + verify(dataService).unregisterReceiver(activeBundle.getId()); + } + + @Test + public void testOutputEmbeddedInAnAlreadyCompletedResponseIsDelivered() throws Exception { + CompletableFuture processBundleResponseFuture = new CompletableFuture<>(); + when(fnApiControlClient.handle(any(BeamFnApi.InstructionRequest.class))) + .thenReturn(processBundleResponseFuture); + FullWindowedValueCoder coder = + FullWindowedValueCoder.of(StringUtf8Coder.of(), Coder.INSTANCE); + BundleProcessor processor = + sdkHarnessClient.getProcessor( + descriptor, + Collections.singletonList( + RemoteInputDestination.of( + (FullWindowedValueCoder) coder, SDK_GRPC_READ_TRANSFORM))); + when(dataService.createOutboundAggregator(any(), anyBoolean())) + .thenReturn(mock(BeamFnDataOutboundAggregator.class)); + + List received = Collections.synchronizedList(new ArrayList<>()); + try (ActiveBundle activeBundle = + processor.newBundle( + ImmutableMap.of( + SDK_GRPC_WRITE_TRANSFORM, + RemoteOutputReceiver.of(ByteArrayCoder.of(), received::add)), + BundleProgressHandler.ignored())) { + Elements embedded = + outputMessages(activeBundle.getId(), 1).get(0).toBuilder() + .mergeFrom(lastOutputMessage(activeBundle.getId())) + .build(); + processBundleResponseFuture.complete( + InstructionResponse.newBuilder() + .setProcessBundle(ProcessBundleResponse.newBuilder().setElements(embedded)) + .build()); + } + assertEquals(1, received.size()); + } + + @Test + public void testFailedOutputDrainDoesNotReportCompletion() throws Exception { + // The SDK finishes the bundle normally, but consuming its output fails. The bundle has + // failed, so it must not report completion, register residuals, or request finalization. + CompletableFuture processBundleResponseFuture = new CompletableFuture<>(); + when(fnApiControlClient.handle(any(BeamFnApi.InstructionRequest.class))) + .thenReturn(processBundleResponseFuture); + FullWindowedValueCoder coder = + FullWindowedValueCoder.of(StringUtf8Coder.of(), Coder.INSTANCE); + BundleProcessor processor = + sdkHarnessClient.getProcessor( + descriptor, + Collections.singletonList( + RemoteInputDestination.of( + (FullWindowedValueCoder) coder, SDK_GRPC_READ_TRANSFORM))); + doNothing().when(dataService).registerReceiver(any(), outputReceiverCaptor.capture()); + when(dataService.createOutboundAggregator(any(), anyBoolean())) + .thenReturn(mock(BeamFnDataOutboundAggregator.class)); + + BundleProgressHandler mockProgressHandler = mock(BundleProgressHandler.class); + BundleCheckpointHandler mockCheckpointHandler = mock(BundleCheckpointHandler.class); + BundleFinalizationHandler mockFinalizationHandler = mock(BundleFinalizationHandler.class); + RuntimeException receiverFailure = new RuntimeException("receiver failed"); + + ActiveBundle activeBundle = + processor.newBundle( + ImmutableMap.of( + SDK_GRPC_WRITE_TRANSFORM, + RemoteOutputReceiver.of( + ByteArrayCoder.of(), + value -> { + throw receiverFailure; + })), + Collections.emptyMap(), + (request) -> { + throw new UnsupportedOperationException(); + }, + mockProgressHandler, + mock(BundleSplitHandler.class), + mockCheckpointHandler, + mockFinalizationHandler); + outputReceiverCaptor.getValue().accept(outputMessages(activeBundle.getId(), 1).get(0)); + processBundleResponseFuture.complete( + InstructionResponse.newBuilder() + .setProcessBundle( + ProcessBundleResponse.newBuilder() + .setRequiresFinalization(true) + .addResidualRoots(DelayedBundleApplication.getDefaultInstance())) + .build()); + + Exception thrown = assertThrows(Exception.class, activeBundle::close); + assertSame(receiverFailure, thrown); + verify(mockProgressHandler, never()).onCompleted(any()); + verifyNoMoreInteractions(mockCheckpointHandler, mockFinalizationHandler); + } + private static class TestFn extends DoFn { @ProcessElement public void processElement(ProcessContext context) {