Simplify budget distribution logic and new worker metadata consumption by m-trieu · Pull Request #32775 · apache/beam

m-trieu · 2024-10-14T22:18:33Z

Remove the budget distributor thread and move the managing of each server's internal budget to the direct GetWorkStream.

R: @arunpandianp @scwhittle

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

Mention the appropriate issue in your description (for example: addresses #123), if applicable. This will automatically add a link to the pull request in the issue. If you would like the issue to automatically close on merging the pull request, comment fixes #<ISSUE NUMBER> instead.
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 or the workflows README to see a list of phrases to trigger workflows.

github-actions · 2024-10-14T23:05:55Z

Assigning reviewers. If you would like to opt out of this review, comment assign to next reviewer:

R: @shunping added as fallback since no labels match configuration

Available commands:

stop reviewer notifications - opt out of the automated review tooling
remind me after tests pass - tag the comment author after tests pass
waiting on author - shift the attention set back to the author (any comment or push by the author will return the attention set to the reviewers)

The PR bot will only process comments in the main thread (not review comments).

...n/java/org/apache/beam/runners/dataflow/worker/streaming/harness/GlobalDataStreamSender.java

arunpandianp · 2024-10-15T06:21:33Z

...ain/java/org/apache/beam/runners/dataflow/worker/streaming/harness/WindmillStreamSender.java

Since started is set inside startStreams after getWorkStream.get() without mutexes guarding them, started can be false here setBudget is called by a different thread when startStreams is in the middle of execution. do we need to change the synchronization in this class?

that is in this PR #32774 and i added synchronization there
this is still the memoized Supplier.get() which is threadsafe

arunpandianp · 2024-10-15T06:32:50Z

...va/org/apache/beam/runners/dataflow/worker/windmill/client/grpc/GrpcDirectGetWorkStream.java

Suggested change

private void sendRequestExtension(GetWorkBudget extension) {

private void maybeSendRequestExtension(GetWorkBudget extension) {

scwhittle · 2024-10-15T11:13:24Z

...pache/beam/runners/dataflow/worker/streaming/harness/FanOutStreamingEngineWorkerHarness.java

nit: just do = new Object() here?

...n/java/org/apache/beam/runners/dataflow/worker/streaming/harness/GlobalDataStreamSender.java

…ream, simplify worker metadata consumption

scwhittle · 2024-10-16T09:27:04Z

...pache/beam/runners/dataflow/worker/streaming/harness/FanOutStreamingEngineWorkerHarness.java

@@ -219,191 +201,195 @@ static FanOutStreamingEngineWorkerHarness forTesting(
  @SuppressWarnings("ReturnValueIgnored")


think this can be removed after removing get()

scwhittle · 2024-10-16T09:33:49Z

...pache/beam/runners/dataflow/worker/streaming/harness/FanOutStreamingEngineWorkerHarness.java

+  @SuppressWarnings("FutureReturnValueIgnored")
+  private void closeStaleStreams(WindmillEndpoints newWindmillEndpoints) {
+    StreamingEngineBackends currentBackends = backends.get();
+    ImmutableMap<Endpoint, WindmillStreamSender> currentWindmillStreams =


remove the variable and just do currentBackends.windmillStreams().entrySet() ... below

scwhittle · 2024-10-16T09:36:05Z

...pache/beam/runners/dataflow/worker/streaming/harness/FanOutStreamingEngineWorkerHarness.java

-    newWorkerMetadataConsumer.shutdownNow();
+    Preconditions.checkState(started, "FanOutStreamingEngineWorkerHarness never started.");
+    Preconditions.checkNotNull(getWorkerMetadataStream).halfClose();
+    workerMetadataConsumer.shutdownNow();


should this shutdown all the stream senders (perhaps could call closeStaleStreams(emptyBackends))

scwhittle · 2024-10-16T09:36:47Z

...pache/beam/runners/dataflow/worker/streaming/harness/FanOutStreamingEngineWorkerHarness.java

-    newWorkerMetadataPublisher.shutdownNow();
-    newWorkerMetadataConsumer.shutdownNow();
+    Preconditions.checkState(started, "FanOutStreamingEngineWorkerHarness never started.");
+    Preconditions.checkNotNull(getWorkerMetadataStream).halfClose();


seems like we should abort/close the stream below if it doesn't terminate within some amount of time after half-closing.

Done

Used shutdown instead, and then used awaitTermination

scwhittle · 2024-10-16T09:37:58Z

...pache/beam/runners/dataflow/worker/streaming/harness/FanOutStreamingEngineWorkerHarness.java


  @VisibleForTesting
  @Override
  public synchronized void shutdown() {


should we also shutdown the windmillStreamManager? (after possibly closing streams below)

scwhittle · 2024-10-16T10:31:37Z

...er/src/main/java/org/apache/beam/runners/dataflow/worker/windmill/client/WindmillStream.java


-    /** Returns the remaining in-flight {@link GetWorkBudget}. */
-    GetWorkBudget remainingBudget();
+    default void setBudget(GetWorkBudget newBudget) {


should this be the virtual method and the long variant defaulted to making a budget?
The implementation makes a budget object

scwhittle · 2024-10-16T10:35:19Z

...va/org/apache/beam/runners/dataflow/worker/windmill/client/grpc/GrpcDirectGetWorkStream.java

+          new AtomicLong(), new AtomicLong(), new AtomicLong(), new AtomicLong());
+    }
+
+    abstract AtomicLong itemsRequested();


how about just using synchronized instead of lots of separate atomics? Multiple atomic ops might be worse performance anyway and it means we might have weird races where they are inconsistently updated.

Done also added tests for this class

can the members be changed to just raw longs/objects? The accessors just need to be synchronized as well.

Seems like this could be easier without autovalue since we don't need the accessors eather.

scwhittle · 2024-10-16T10:36:08Z

...va/org/apache/beam/runners/dataflow/worker/windmill/client/grpc/GrpcDirectGetWorkStream.java

+    this.heartbeatSender = heartbeatSender;
+    this.workCommitter = workCommitter;
+    this.getDataClient = getDataClient;
+    this.maxGetWorkBudget =


can this max be moved into the tracker and synchronized within it?
As is there are races between setting and using this when calling into the tracker.

scwhittle · 2024-10-16T10:42:56Z

...g/apache/beam/runners/dataflow/worker/windmill/client/grpc/stubs/WindmillChannelFactory.java

      default:
        throw new UnsupportedOperationException(
-            "Only IPV6, GCP_SERVICE_ADDRESS, AUTHENTICATED_GCP_SERVICE_ADDRESS are supported WindmillServiceAddresses.");
+            "Only IPV6, GCP_SERVICE_ADDRESS, AUTHENTICATED_GCP_SERVICE_ADDRESS are supported"


remove IPV6 from comment

scwhittle · 2024-10-16T10:45:40Z

...g/apache/beam/runners/dataflow/worker/windmill/work/budget/EvenGetWorkBudgetDistributor.java

      ImmutableCollection<T> streams, GetWorkBudget totalGetWorkBudget) {
-    GetWorkBudget activeWorkBudget = activeWorkBudgetSupplier.get();
-    LOG.info("Current active work budget: {}", activeWorkBudget);
    // TODO: Fix possibly non-deterministic handing out of budgets.


remove, doesn't seem like the same if we are not doing deltas

…tWorkStream

m-trieu · 2024-10-17T13:06:50Z

back to you @scwhittle
Thanks!

scwhittle · 2024-10-17T08:14:29Z

...pache/beam/runners/dataflow/worker/streaming/harness/FanOutStreamingEngineWorkerHarness.java

-                    () -> closeStreamSender(entry.getKey(), entry.getValue()),
-                    windmillStreamManager));
+            entry -> {
+              CompletableFuture<Void> ignored =


is this any different than just executing directly? if not it seems simpler to avoid the future.

windmillStreamManager.execute(
() -> closeStreamSender(entry.getKey(), entry.getValue()))

scwhittle · 2024-10-17T08:14:55Z

...pache/beam/runners/dataflow/worker/streaming/harness/FanOutStreamingEngineWorkerHarness.java

-                CompletableFuture.runAsync(
-                    () -> closeStreamSender(sender.endpoint(), sender), windmillStreamManager));
+            sender -> {
+              CompletableFuture<Void> ignored =


scwhittle · 2024-10-17T08:16:31Z

...worker/src/main/java/org/apache/beam/runners/dataflow/worker/windmill/WindmillEndpoints.java


+  public static WindmillEndpoints none() {
+    return WindmillEndpoints.builder()
+        .setVersion(Long.MAX_VALUE)


min seems safer. Otherwise if somehow none() was observed the logic to ensure version is increasing mean's we'd never process another endpoint set

scwhittle · 2024-10-17T08:24:30Z

...va/org/apache/beam/runners/dataflow/worker/windmill/client/grpc/GrpcDirectGetWorkStream.java

            + "last sent request: %s. ",
        workItemAssemblers.size(),
-        maxGetWorkBudget.get(),
+        budgetTracker.maxGetWorkBudget().get(),


could move html generation into budgettracker and not need all the accessors. If we change how the tracker works in the future we might want to show more too.

scwhittle · 2024-10-17T08:24:35Z

...va/org/apache/beam/runners/dataflow/worker/windmill/client/grpc/GrpcDirectGetWorkStream.java

+          new AtomicLong(), new AtomicLong(), new AtomicLong(), new AtomicLong());
+    }
+
+    abstract AtomicLong itemsRequested();


can the members be changed to just raw longs/objects? The accessors just need to be synchronized as well.

Seems like this could be easier without autovalue since we don't need the accessors eather.

scwhittle · 2024-10-18T13:04:19Z

...rg/apache/beam/runners/dataflow/worker/windmill/client/grpc/GrpcDirectGetWorkStreamTest.java

+    GetWorkStreamTestStub testStub = new GetWorkStreamTestStub(requestObserver);
+    stream = createGetWorkStream(testStub, GetWorkBudget.noBudget(), throttleTimer);
+    stream.startThrottleTimer();
+    testStub.injectResponse(Windmill.StreamingGetWorkResponseChunk.getDefaultInstance());


does this run inline? otherwise it seems like it coudl be racy below that unthrottling happens?

yea it does, this is the actual response observer we create in the AbstractWindmillStream

scwhittle · 2024-10-18T13:05:15Z

...rg/apache/beam/runners/dataflow/worker/windmill/client/grpc/GrpcDirectGetWorkStreamTest.java

+
+  private static class GetWorkStreamTestStub
+      extends CloudWindmillServiceV1Alpha1Grpc.CloudWindmillServiceV1Alpha1ImplBase {
+    private final TestGetWorkRequestObserver requestObserver;


might as well mark volatile to prevent races

made the mutable state in TestGetWorkRequestObserver threadsafe, to prevent races
this field is sent once per lifetime of the stream so final is ok i think

scwhittle · 2024-10-18T13:06:17Z

...ache/beam/runners/dataflow/worker/windmill/work/budget/EvenGetWorkBudgetDistributorTest.java

-    GetWorkBudgetSpender getWorkBudgetSpender =
-        spy(createGetWorkBudgetOwnerWithRemainingBudgetOf(GetWorkBudget.noBudget()));
-    createBudgetDistributor(1L)
+    GetWorkBudgetSpender getWorkBudgetSpender = spy(createGetWorkBudgetOwner());


remove spy here? already done in the helper method

ditto for below

scwhittle · 2024-10-18T13:08:16Z

...ache/beam/runners/dataflow/worker/windmill/work/budget/EvenGetWorkBudgetDistributorTest.java

-
  @Test
  public void testDistributeBudget_distributesBudgetEvenlyIfPossible() {
    long totalItemsAndBytes = 10L;


would be better to have different items and bytes values to confirm distributor doesn't mix them up internally

scwhittle · 2024-10-18T13:08:50Z

...ache/beam/runners/dataflow/worker/windmill/work/budget/EvenGetWorkBudgetDistributorTest.java

@@ -192,17 +82,17 @@ public void testDistributeBudget_distributesBudgetEvenlyIfPossible() {
    streams.forEach(


just skip the math in the test and inline the right values?
The math is just copying what we have in the impl, if there is some bug in the impl hard coding the values at least is a sanity check.

ditto below.

m-trieu · 2024-10-21T03:19:21Z

thanks @scwhittle !
back to you

scwhittle · 2024-10-21T08:07:01Z

Run Java PreCommit

scwhittle · 2024-10-21T07:56:51Z

...va/org/apache/beam/runners/dataflow/worker/windmill/client/grpc/GrpcDirectGetWorkStream.java

-    }
+  private static final class GetWorkBudgetTracker {
+
+    @GuardedBy("GetWorkBudgetTracker.this")


you can just use "this"
https://javadoc.io/doc/com.google.code.findbugs/annotations/latest/javax/annotation/concurrent/GuardedBy.html

Abacn · 2024-11-07T03:05:52Z

It seems "FanOutStreamingEngineWorkerHarnessTest" becomes increasingly flaky recently, e.g. on schedule run:

https://github.com/apache/beam/runs/32620998299

and on a PR run:

https://github.com/apache/beam/runs/32622849052

error:

java.lang.AssertionError
	at org.junit.Assert.fail(Assert.java:87)
	at org.junit.Assert.assertTrue(Assert.java:42)
	at org.junit.Assert.assertTrue(Assert.java:53)
	at org.apache.beam.runners.dataflow.worker.streaming.harness.FanOutStreamingEngineWorkerHarnessTest.testStreamsStartCorrectly(FanOutStreamingEngineWorkerHarnessTest.java:217)

github-actions bot added runners dataflow labels Oct 14, 2024

github-actions bot added the Next Action: Reviewers label Oct 14, 2024

arunpandianp reviewed Oct 15, 2024

View reviewed changes

scwhittle requested changes Oct 15, 2024

View reviewed changes

m-trieu mentioned this pull request Oct 15, 2024

add shutdown and start mechanics to windmill streams #32774

Merged

3 tasks

simplify budget distribution logic moving it into GrpcDirectGetWorkSt…

b866d42

…ream, simplify worker metadata consumption

m-trieu force-pushed the mt-simplify-budget branch from 759ca44 to 53af347 Compare October 15, 2024 17:49

address PR comments

496d1ed

m-trieu force-pushed the mt-simplify-budget branch from 53af347 to 496d1ed Compare October 15, 2024 18:35

scwhittle requested changes Oct 16, 2024

View reviewed changes

address PR comments, add unit tests around budgeting for GrpcDirectGe…

3f6888d

…tWorkStream

scwhittle requested changes Oct 18, 2024

View reviewed changes

address pr comments

e62c55a

scwhittle approved these changes Oct 21, 2024

View reviewed changes

scwhittle merged commit 68f1543 into apache:master Oct 21, 2024

Abacn mentioned this pull request Nov 7, 2024

Remove antlr from dependency from java-core pom #33030

Merged

3 tasks

	private void sendRequestExtension(GetWorkBudget extension) {
	private void maybeSendRequestExtension(GetWorkBudget extension) {

		@@ -219,191 +201,195 @@ static FanOutStreamingEngineWorkerHarness forTesting(
		@SuppressWarnings("ReturnValueIgnored")

		@@ -192,17 +82,17 @@ public void testDistributeBudget_distributesBudgetEvenlyIfPossible() {
		streams.forEach(

Conversation

m-trieu commented Oct 14, 2024

GitHub Actions Tests Status (on master branch)

Uh oh!

github-actions bot commented Oct 14, 2024

Uh oh!

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

m-trieu commented Oct 17, 2024

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment