Vortex performance improvement: Enable multiple stream clients per worker #17550

prodriguezdefino · 2022-05-04T22:30:28Z

Changing the StorageWrite stream append client cache to use a list of entries instead of just one entry per stream.

In the majority of the cases the default configuration would be sufficient to address the ingestion volume, but there are cases were the ingestion will need a higher level of parallelism. Setting --numStorageWriteApiStreams=X to a value > 1 will help to increase the ingestion parallelism level, users should be mindful of changing this number to a high one since the number of open connections will rapidly scale with the number of workers potentially reaching BigQuery quota/limits (see: https://cloud.google.com/bigquery/quotas#write-api-limits).

Please add a meaningful description for your change here

Thank you for your contribution! Follow this checklist to help us incorporate your contribution quickly and easily:

Choose reviewer(s) and mention them in a comment (R: @username).
Format the pull request title like [BEAM-XXX] Fixes bug in ApproximateQuantiles, where you replace BEAM-XXX with the appropriate JIRA issue, if applicable. This will automatically link the pull request to the issue.
Update CHANGES.md with noteworthy changes.
If this contribution is large, please file an Apache Individual Contributor License Agreement.

See the Contributor Guide for more tips on how to make review process smoother.

To check the build health, please visit https://github.com/apache/beam/blob/master/.test-infra/BUILD_STATUS.md

GitHub Actions Tests Status (on master branch)

See CI.md for more information about GitHub Actions CI.

…es instead of just one.

asf-ci · 2022-05-04T22:30:30Z

Can one of the admins verify this patch?

asf-ci · 2022-05-04T22:30:30Z

Can one of the admins verify this patch?

prodriguezdefino · 2022-05-04T22:31:22Z

R: @reuvenlax

prodriguezdefino · 2022-05-05T04:55:28Z

Run Java PreCommit

...google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/BigQueryOptions.java

reuvenlax · 2022-05-04T22:59:21Z

...tform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/StorageApiWriteUnshardedRecords.java

@@ -86,15 +89,18 @@
  // (any access of the cache could trigger element expiration). Therefore most used of
  // APPEND_CLIENTS should
  // synchronize.
-  private static final Cache<String, StreamAppendClient> APPEND_CLIENTS =
+  private static final Cache<String, List<StreamAppendClient>> APPEND_CLIENTS =


what did you find about the cost of synchronization?

reuvenlax · 2022-05-04T23:00:21Z

...tform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/StorageApiWriteUnshardedRecords.java

@@ -129,6 +135,9 @@ public StorageApiWriteUnshardedRecords(
  public PCollection<Void> expand(PCollection<KV<DestinationT, StorageApiWritePayload>> input) {
    String operationName = input.getName() + "/" + getName();
    BigQueryOptions options = input.getPipeline().getOptions().as(BigQueryOptions.class);
+    // default value from options is 0, so we set at least one client
+    Integer numStreams =
+        options.getNumStorageWriteApiStreams() == 0 ? 1 : options.getNumStorageWriteApiStreams();


create a new option that defaults to 1

reuvenlax · 2022-05-04T23:08:17Z

...tform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/StorageApiWriteUnshardedRecords.java

-                        () ->
-                            datasetService.getStreamAppendClient(
-                                streamName, descriptorWrapper.descriptor));
+                    APPEND_CLIENTS.get(streamName, () -> generateClients()).get(clientNumber);


Instead of generating all clients eagerly, let's do it lazily. Initialize a List with count copies of Optional.empty(). Then do
this.streamAppendCient = APPEND_CLIENTS.get(streamName, this.generateClients).get(clientNumber).get().orElseGet(this.getStreamAppendClient).

FYI you could also do this with null if you don't care to use Optional here.

Instead of making the client creation lazy, I reverted the cache structure back to have a single client per entry. But now, the cache key is the stream name + the assigned client number.

reuvenlax · 2022-05-05T17:05:53Z

...tform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/StorageApiWriteUnshardedRecords.java

-                && System.identityHashCode(cachedAppendClient)
+            List<StreamAppendClient> cachedAppendClients = APPEND_CLIENTS.getIfPresent(streamName);
+            if (cachedAppendClients != null
+                && System.identityHashCode(cachedAppendClients.get(clientNumber))
                    == System.identityHashCode(streamAppendClient)) {


This isn't quite right - we're now invalidating all of the StreamWriters when any one of them fails. I think instead you want to just null out the one that failed and allow it to be recreated the next get.

The invalidation here corresponds to a schema mismatch, shouldn't all the clients be invalidated for a particular stream?

made the changes to only invalidate the writer in use by the bundle, not all of them.

reuvenlax · 2022-05-05T17:06:07Z

...tform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/StorageApiWriteUnshardedRecords.java

    private final StorageApiDynamicDestinations<ElementT, DestinationT> dynamicDestinations;
    private final BigQueryServices bqServices;
    private final boolean useDefaultStream;
+    // default append client count to 1
+    private Integer streamAppendClientCount = 1;


why not private int?

… eviction logic

reuvenlax · 2022-05-11T20:27:32Z

...google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/BigQueryOptions.java

+      "The number of stream append clients indicated will be allocated at a per worker and destination "
+          + "basis. A large value can cause a large pipeline to go over the BigQuery connection quota quickly. "
+          + "With low-mid volume pipelines using the default configuration should be enough.")
+  @Default.Integer(1)


A bit confusing - need to clarify that this only applies for at-least once writes using the default stream

reuvenlax · 2022-05-11T20:28:31Z

...rm/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/StorageApiWriteRecordsInconsistent.java

@@ -50,14 +49,16 @@ public PCollection<Void> expand(PCollection<KV<DestinationT, StorageApiWritePayl
    BigQueryOptions bigQueryOptions = input.getPipeline().getOptions().as(BigQueryOptions.class);
    // Append records to the Storage API streams.
    input.apply(
-        "Write Records",


Changing transform names can affect update compatibility - do you need this?

No I don't, but this does not follow the same convention other apply labels use (no spaces on names). Should I revert?

reuvenlax · 2022-05-11T20:38:59Z

...tform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/StorageApiWriteUnshardedRecords.java

@@ -197,6 +204,10 @@ String getDefaultStreamName() {
        return BigQueryHelpers.stripPartitionDecorator(tableUrn) + "/streams/_default";
      }

+      String getStreamAppendClientCacheEntryName() {
+        return getDefaultStreamName() + "-client" + clientNumber;
+      }


This is a bit weird, since this code doesn't always use the default stream. Now the cache is probably not needed in the non default stream case (since we'll create a new stream for every bundle), however if we change that we need to rename the cache and also make sure to close the client (since right now we rely on the cache removal listener to close the client)

I will make the changes to consider default streams and per bundle on-demand streams.

reuvenlax · 2022-05-11T20:40:23Z

...tform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/StorageApiWriteUnshardedRecords.java

-                        () ->
-                            datasetService.getStreamAppendClient(
-                                streamName, descriptorWrapper.descriptor));
+                        getStreamAppendClientCacheEntryName(), () -> createStreamAppendClient());


getStreamAppendClientCacheEntryName doesn't necessarily return the stream name we used

It will also be broken if you have two sinks in the pipeline, one using the default stream and one not.

I can see the problem of mixing the default with dynamically gen streams, I will change that.

…eanup/init to teardown/setup when using default stream

reuvenlax · 2022-05-18T18:25:24Z

...tform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/StorageApiWriteUnshardedRecords.java

@@ -409,9 +447,18 @@ private void initializeDatasetService(PipelineOptions pipelineOptions) {
      }
    }

+    @Setup
+    public void setup() {


I'm not sure what this is adding?

I've seen evidence that closing one stream append client for the _default stream causes a cascade close of the other ones, so moving the state for destinations to be reused between bundle executions at least decreased the occurrences of those cascading closes.

I can revert this if this is not the right idea to try.

reuvenlax · 2022-05-18T18:28:30Z

...tform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/StorageApiWriteUnshardedRecords.java

    @StartBundle
    public void startBundle() throws IOException {
-      destinations = Maps.newHashMap();
+      if (!useDefaultStream) {


Why add this if?

Only want to reuse the destination state for the default stream clients.

reuvenlax · 2022-05-18T18:40:45Z

...tform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/StorageApiWriteUnshardedRecords.java

@@ -455,21 +506,26 @@ public void process(
    @FinishBundle
    public void finishBundle(FinishBundleContext context) throws Exception {
      flushAll();
-      for (DestinationState state : destinations.values()) {
-        if (!useDefaultStream) {
+      if (!useDefaultStream) {


Why this change? Won't this prevent us from unpinning the client?

it will, but those clients can be reused across bundle executions.

the clients can be reused regardless - unpinning won't close the client unless close is also called (i.e. if we hit the cache idle timeout)

reuvenlax · 2022-05-18T18:41:52Z

...tform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/StorageApiWriteUnshardedRecords.java

+      if (destinations != null) {
+        for (DestinationState state : destinations.values()) {
+          state.teardown();
+        }


We shouldn't be doing this in teardown

reuvenlax · 2022-05-18T18:42:43Z

...google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/BigQueryOptions.java

  @Default.Integer(0)
  Integer getNumStorageWriteApiStreams();

  void setNumStorageWriteApiStreams(Integer value);

+  @Description(
+      "When using the \"_default\" table stream, this option sets the number of stream append clients that will be allocated "


reference at-least once writes. Users don't know about default streams, as that's an implementation detail

y1chi · 2022-05-19T00:56:41Z

Run Java PreCommit

reuvenlax · 2022-05-19T12:29:45Z

A couple of very minor nits, but otherwise LGTM

reuvenlax · 2022-05-19T12:27:22Z

...tform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/StorageApiWriteUnshardedRecords.java

@@ -213,6 +226,14 @@ String createStreamIfNeeded() {
        return this.streamName;
      }

+      StreamAppendClient generateClient() {


Minor nit - remove try/catch since the calling function already catches the exception

reuvenlax · 2022-05-19T12:29:00Z

...tform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/StorageApiWriteUnshardedRecords.java

@@ -263,11 +281,12 @@ void invalidateWriteStream() {
            // thread has already invalidated
            // and recreated the stream).
            @Nullable
-            StreamAppendClient cachedAppendClient = APPEND_CLIENTS.getIfPresent(streamName);
+            StreamAppendClient cachedAppendClient =
+                APPEND_CLIENTS.getIfPresent(getStreamAppendClientCacheEntryName());


nit - store cachedEntryName in a local variable instead of recomputing twice here

prodriguezdefino · 2022-05-19T15:20:25Z

Run Java PreCommit

reuvenlax · 2022-05-19T16:12:45Z

Run Java PreCommit

y1chi · 2022-05-19T17:24:08Z

Run Java PreCommit

…rker (apache#17550)

…rker (#17550) (#17718) Co-authored-by: pablo rodriguez defino <[email protected]>

changing the StorageWrite stream append client to use a list of entri…

a47e7ee

…es instead of just one.

github-actions bot added gcp io java labels May 4, 2022

ran spotlessApply

a9d9207

reuvenlax reviewed May 5, 2022

View reviewed changes

prodriguezdefino added 2 commits May 9, 2022 15:52

storing single client per cache entry, simplying check, retrieval and…

615373c

… eviction logic

name changed on stage, making more clear the impl selection

8957d66

reuvenlax reviewed May 11, 2022

View reviewed changes

prodriguezdefino added 2 commits May 13, 2022 18:04

reverting to eagerly initialized client streams, moved destination cl…

bf41d95

…eanup/init to teardown/setup when using default stream

re instated per stream cache entry

e35bdc2

prodriguezdefino requested a review from reuvenlax May 17, 2022 00:21

reuvenlax reviewed May 18, 2022

View reviewed changes

prodriguezdefino added 2 commits May 18, 2022 12:15

changing javadoc on Options configuration per review

aa03e4a

reverting the stream append client pin/unpin lifecycle changes

569ea58

reuvenlax reviewed May 19, 2022

View reviewed changes

addressing comments

50664b5

y1chi merged commit 47d8bce into apache:master May 19, 2022

y1chi pushed a commit to y1chi/beam that referenced this pull request May 19, 2022

Vortex performance improvement: Enable multiple stream clients per wo…

9dfd79c

…rker (apache#17550)

y1chi mentioned this pull request May 19, 2022

Vortex performance improvement: Enable multiple stream clients per wo… #17718

Merged

4 tasks

y1chi added a commit that referenced this pull request May 19, 2022

Vortex performance improvement: Enable multiple stream clients per wo…

15a211e

…rker (#17550) (#17718) Co-authored-by: pablo rodriguez defino <[email protected]>

Vortex performance improvement: Enable multiple stream clients per worker #17550

Vortex performance improvement: Enable multiple stream clients per worker #17550

Conversation

prodriguezdefino commented May 4, 2022 • edited Loading

GitHub Actions Tests Status (on master branch)

asf-ci commented May 4, 2022

asf-ci commented May 4, 2022

prodriguezdefino commented May 4, 2022

prodriguezdefino commented May 5, 2022

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

y1chi commented May 19, 2022

reuvenlax commented May 19, 2022

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

prodriguezdefino commented May 19, 2022

reuvenlax commented May 19, 2022

y1chi commented May 19, 2022

prodriguezdefino commented May 4, 2022 •

edited

Loading