Update windmill proto definition #30046

m-trieu · 2024-01-19T02:30:51Z

Internal and external protos were not matching leading to gRPC errors.

make both protos match
add client id to the job headers for communication with Windmill Service
let the client id be owned by GrpcWindmillServer (WindmillServerStub) for now since its injected.
Will need to find a way to inject the client id (maybe from StreamingDataflowWorkerOptions) as well as other components. I wonder if that will be cleaner

r: @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-01-19T03:06:09Z

Checks are failing. Will not request review until checks are succeeding. If you'd like to override that behavior, comment assign set of reviewers

m-trieu · 2024-01-19T03:35:16Z

assign set of reviewers

m-trieu · 2024-01-19T03:35:26Z

run java precommit

github-actions · 2024-01-19T03:36:24Z

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

R: @lostluck 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).

runners/google-cloud-dataflow-java/worker/windmill/src/main/proto/windmill.proto

scwhittle · 2024-01-19T09:55:32Z

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

+
+    return directEndpointIpV6Address.isPresent()
+        ? directEndpointIpV6Address
+        : tryParseEndpointIntoHostAndPort(directEndpoint).map(WindmillServiceAddress::create);


nit, the WindmillServiceAddress::create is common,
so maybe cleaner to do
return
tryParseDirectEndpointIntoIpV6Address(directEndpoint).orElseGet(tryParseEndpointIntoHostAndPort(directEndpoint)).map(WindmillServiceAddress::create);

i am wondering if we should jsut drop the ipv6
since HostAndPort supports ipv6

its possible that the tryParseEndpointIntoHostAndPort fails so it returns Optional and orElseGet requires a signature that returns HostAndPort

scwhittle · 2024-01-19T09:56:54Z

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

    public static Endpoint.Builder builder() {
      return new AutoValue_WindmillEndpoints_Endpoint.Builder();
    }

-    public static Endpoint from(Windmill.WorkerMetadataResponse.Endpoint endpointProto) {
+    public static Endpoint from(
+        Windmill.WorkerMetadataResponse.Endpoint endpointProto, String authenticatingService) {
      Endpoint.Builder endpointBuilder = Endpoint.builder();
      if (endpointProto.hasDirectEndpoint() && !endpointProto.getDirectEndpoint().isEmpty()) {


remove the has check, can rely on default value being empty

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

scwhittle · 2024-01-19T09:59:26Z

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

@@ -76,6 +76,10 @@ public abstract GetWorkStream getWorkStream(
  /** Returns the amount of time the server has been throttled and resets the time to 0. */
  public abstract long getAndResetThrottleTime();

+  public long clientId() {
+    return 0L;


client id is supposed to change on process restart. It is otherwise not a value that has to be coordinated in advance with the service so it can just be a random value as it was before. I think that it can stay in StreamingDataflowWorker because it also doesn't need to be different for each endpoint.

This helps the backend speed up retries, ensure caching is correct etc.

ah ok out of curiosity wouldn't the other components restart also?
looks like we create a new worker in StreamingDataflowWorker.main

And inject the WindmillServerStub (GrpcWindmillServer) via the options.

scwhittle · 2024-01-19T10:01:46Z

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

@@ -45,92 +48,156 @@ class GrpcDispatcherClient {
  private final WindmillStubFactory windmillStubFactory;

  @GuardedBy("this")


seems like this wouldn't need guarding if it is atomic.

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

lostluck · 2024-01-19T19:38:51Z

R: @scwhittle

(Just to get the bot off the hook. The "R" needs to be R and not r.)

github-actions · 2024-01-19T19:40:22Z

Stopping reviewer notifications for this pull request: review requested by someone other than the bot, ceding control

m-trieu · 2024-01-19T22:50:00Z

R: @scwhittle

(Just to get the bot off the hook. The "R" needs to be R and not r.)

ack thanks!

scwhittle · 2024-01-29T12:49:16Z

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

+  }
+
+  private synchronized <T> T randomlySelectNextStub(List<T> stubs) {
+    return stubs.get(rand.nextInt(dispatcherStubs.get().size()));


mixing stubs and dispatcherStubs

changed to use stubs.size()

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

scwhittle · 2024-01-29T12:54:51Z

...va/worker/src/main/java/org/apache/beam/runners/dataflow/worker/StreamingDataflowWorker.java

+        !options.isEnableStreamingEngine() && options.getLocalWindmillHostport() != null
+            ? GrpcWindmillServer.LOCALHOST_BACKOFF
+            : GrpcWindmillServer.MAX_BACKOFF;
+    GrpcWindmillStreamFactory windmillStreamFactory =


seems odd to create this here if the server isn't grpc server.

Can we just pass in somethign generic like JobHeader to windmillServer.start?

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

scwhittle · 2024-02-05T09:56:52Z

...va/worker/src/main/java/org/apache/beam/runners/dataflow/worker/StreamingDataflowWorker.java

    this.windmillServer = options.getWindmillServerStub();
    this.windmillServer.setProcessHeartbeatResponses(this::handleHeartbeatResponses);
+    windmillServer.start(


could we avoid the start method (which is more complex than a factory method taking necessary requirements and initalizing) by moving the grpc server initialization into this file?

It's kind of odd that it is within getWindmillServerStub default factory. It appears that is done so that the test can inject a fake server. One idea would be to remove the default factory from options and create a server here if the option is null but use injected server otherwise.

Then you could just create the GrpcWindmillServer directly with either a factory or constructor taking the JobHeader without modifying the WindmillServerStub interface.

changed
opted to allow StreamingDataflowWorker to inject a stream factory

scwhittle · 2024-02-05T09:59:46Z

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

  }

-  synchronized boolean isReady() {
-    return !dispatcherStubs.isEmpty();
+  boolean isReady() {


this seems like it could lead to racy usage
if (client.isReady()) {
client.xxx
}
but since the dispatcherStubs can change in-between, there may not be a stub by the time the next call is made.

Can this be removed?

We use this to make sure that the endpoints are set in StreamingDataflowWorker before we fetch the pipeline config as well as schedule periodic config fetches

beam/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/StreamingDataflowWorker.java

Line 1615 in f08058c

if (!windmillServer.isReady()) {

if we remove this we would need another mechanism for that.

we could use some kind of publish/subscribe mechanism to signal when the endpoint is ready (kind of like a count down latch) and then send another signal during an update (i.e was ready with endpoints, but is no longer ready).

How often does this value change? I would think that the service endpoint itself would not change during pipeline execution. @scwhittle

i guess we can do it async

how often does this value change? @scwhittle

We never really expect it to change. It was added in case we needed to migrate dispatcher endpoint since there are possibly stale long-running client processes.

How about ignoring updates that are empty? Then we only go from not read -> ready and not backwards? That prevents the possible race, shouldn't cause any problems and is simpler than another type of notification.

this is the current behavior which won't allow empty updates

synchronized void consumeWindmillDispatcherEndpoints( ImmutableSet<HostAndPort> dispatcherEndpoints) { ImmutableSet<HostAndPort> currentDispatcherEndpoints = dispatcherStubs.get().dispatcherEndpoints(); Preconditions.checkArgument( dispatcherEndpoints != null && !dispatcherEndpoints.isEmpty(), "Cannot set dispatcher endpoints to nothing."); if (this.dispatcherEndpoints.equals(dispatcherEndpoints)) { if (currentDispatcherEndpoints.equals(dispatcherEndpoints)) { // The endpoints are equal don't recreate the stubs. return; }....

We currently will throw an exception if we try to set the endpoints to nothing, should we just log a warning and ignore the update here?

that seems fine, and we can just comment here that once ready it remains ready

could rename isReady to endpointsInitialized() or something that clarifies what it means more and is clearer it won't go back to false.

...in/java/org/apache/beam/runners/dataflow/worker/windmill/client/grpc/GrpcWindmillServer.java

scwhittle · 2024-02-05T10:04:58Z

...in/java/org/apache/beam/runners/dataflow/worker/windmill/client/grpc/GrpcWindmillServer.java

@@ -101,36 +103,21 @@ public final class GrpcWindmillServer extends WindmillServerStub {
  // newer ComputationHeartbeatRequests.
  private final boolean sendKeyedGetDataRequests;
  private Consumer<List<ComputationHeartbeatResponse>> processHeartbeatResponses;
+  private @Nullable GrpcWindmillStreamFactory windmillStreamFactory;


if we can change constructor to require the streamfactory this nullable and preconditons below can be removed, see other comment

scwhittle · 2024-02-05T10:07:00Z

...he/beam/runners/dataflow/worker/windmill/client/grpc/stubs/InProcessWindmillStubFactory.java

+  private final String testName;
+  private final Function<String, ManagedChannel> channelFactory;
+
+  InProcessWindmillStubFactory(String testName, Function<String, ManagedChannel> channelFactory) {


nit: coudl just take a Supplier
If tests have a Function<String, ManagedChannel> they can always bind the test name themselves.

scwhittle · 2024-02-05T10:10:07Z

...he/beam/runners/dataflow/worker/windmill/client/grpc/stubs/InProcessWindmillStubFactory.java

+import org.apache.beam.vendor.grpc.v1p60p1.io.grpc.ManagedChannel;
+
+/**
+ * Creates in process stubs to talk to Streaming Engine. Only recommended to be used for testing.


should this be moved to test folder and renamed FakeWindmillStubFactory? it ignores the serviceAddress

scwhittle · 2024-02-05T10:11:09Z

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

+import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.annotations.VisibleForTesting;
+
+@Internal
+public final class WindmillStubFactories {


should we get rid of this one? If the fake is moved to test folder this just has single method. I think callers could use RemoteWindmillStubFactory constructor directly

m-trieu · 2024-02-09T06:53:13Z

theres some tests failing for windows that are unrelated (passing in other configurations)

scwhittle · 2024-02-09T09:48:04Z

...va/worker/src/main/java/org/apache/beam/runners/dataflow/worker/StreamingDataflowWorker.java

-
+  public static StreamingDataflowWorker fromOptions(StreamingDataflowWorkerOptions options) {
+    StreamingDataflowWorkerOptions streamingOptions =
+        options.as(StreamingDataflowWorkerOptions.class);


no longer needed, already right option type

scwhittle · 2024-02-09T09:51:55Z

.../worker/src/main/java/org/apache/beam/runners/dataflow/worker/streaming/ActiveWorkState.java

+      ShardedKey shardedKey = entry.getKey();
+      Optional.ofNullable(failedWork.get(shardedKey.shardingKey()))
+          .ifPresent(
+              failedWorkTokens ->


nit: I find this harder to read than just having a variable storing result of failedWork.get and using an if statement, as it's harder to eyeball and takes a little more effort to realize that this is processing on the result of failedWork.get

scwhittle · 2024-02-09T09:54:30Z

.../src/main/java/org/apache/beam/runners/dataflow/worker/streaming/WorkHeartbeatProcessor.java

+  private final Function<String, Optional<ComputationState>> computationStateFetcher;
+
+  public WorkHeartbeatProcessor(
+      Function<String, Optional<ComputationState>> computationStateFetcher) {


comment on what string is

scwhittle · 2024-02-09T09:55:54Z

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

  }

-  synchronized boolean isReady() {
-    return !dispatcherStubs.isEmpty();
+  boolean isReady() {


scwhittle · 2024-02-09T10:08:47Z

...ava/org/apache/beam/runners/dataflow/worker/windmill/client/grpc/GrpcWindmillServerTest.java

  private static final Logger LOG = LoggerFactory.getLogger(GrpcWindmillServerTest.class);
  private static final int STREAM_CHUNK_SIZE = 2 << 20;
+  private final long clientId = new Random().nextLong();


nit: just hard-code something?

scwhittle · 2024-02-13T12:42:12Z

.../worker/src/main/java/org/apache/beam/runners/dataflow/worker/streaming/FailedWorkToken.java

+/*
+ * 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


merge with #29082 and remove this class

scwhittle · 2024-02-13T12:46:27Z

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

  }

-  synchronized boolean isReady() {
-    return !dispatcherStubs.isEmpty();
+  boolean isReady() {


We never really expect it to change. It was added in case we needed to migrate dispatcher endpoint since there are possibly stale long-running client processes.

How about ignoring updates that are empty? Then we only go from not read -> ready and not backwards? That prevents the possible race, shouldn't cause any problems and is simpler than another type of notification.

scwhittle · 2024-02-13T12:50:52Z

...pache/beam/runners/dataflow/worker/windmill/client/grpc/stubs/RemoteWindmillStubFactory.java

+
+  private ManagedChannel createChannel(WindmillServiceAddress serviceAddress) {
+    ManagedChannel channel = remoteChannel(serviceAddress, rpcChannelTimeoutSec);
+    return useIsolatedChannels ? IsolationChannel.create(() -> channel) : channel;


this is using the same ManagedChannel for all the separate channels the IsolationChannel is trying to create.

Instead you can share the code this way. Maybe add a comment since it wasn't clear before

Producer channelFactory = () -> remoteChannel(service_address, rpcChannelTimeoutSec);
return useIsolatedChannels ? IsolationChannel.create(channelFactory) : channelFactory();

scwhittle · 2024-02-13T12:54:52Z

.../src/main/java/org/apache/beam/runners/dataflow/worker/streaming/WorkHeartbeatProcessor.java

+          computationHeartbeatResponse.getHeartbeatResponsesList()) {
+        if (heartbeatResponse.getFailed()) {
+          failedWork
+              .computeIfAbsent(heartbeatResponse.getShardingKey(), key -> new ArrayList<>())


update based upon changes from other PR

…rker

m-trieu · 2024-02-15T00:31:42Z

failures are not related to changes
attempt to fix flaky streaming tests in #30322

ready for another look! @scwhittle looks like we already do not allow empty updates to the dispatcher endpoints

scwhittle

just a couple minor comments, otherwise looks good

scwhittle · 2024-02-15T10:54:12Z

.../src/main/java/org/apache/beam/runners/dataflow/worker/streaming/WorkHeartbeatProcessor.java

+import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.Multimap;
+
+@Internal
+public final class WorkHeartbeatProcessor implements Consumer<List<ComputationHeartbeatResponse>> {


WorkHeartbeatResponseProcessor?
might help prevent confusion that it is processing for outgoing heartbeats

done and added comment

scwhittle · 2024-02-15T11:01:02Z

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

  }

-  synchronized boolean isReady() {
-    return !dispatcherStubs.isEmpty();
+  boolean isReady() {


that seems fine, and we can just comment here that once ready it remains ready

could rename isReady to endpointsInitialized() or something that clarifies what it means more and is clearer it won't go back to false.

github-actions bot added runners dataflow labels Jan 19, 2024

github-actions bot added the Next Action: Reviewers label Jan 19, 2024

scwhittle requested changes Jan 19, 2024

View reviewed changes

m-trieu force-pushed the mt-metadata-api branch 2 times, most recently from 68409ad to 4de69da Compare January 24, 2024 00:20

scwhittle requested changes Jan 29, 2024

View reviewed changes

scwhittle requested changes Feb 5, 2024

View reviewed changes

m-trieu force-pushed the mt-metadata-api branch 3 times, most recently from d04cdf9 to 23d6047 Compare February 8, 2024 21:04

scwhittle requested changes Feb 9, 2024

View reviewed changes

m-trieu force-pushed the mt-metadata-api branch from b9e9949 to a60ec02 Compare February 12, 2024 22:41

scwhittle requested changes Feb 13, 2024

View reviewed changes

m-trieu added 8 commits February 13, 2024 10:14

make external and internal windmill proto defs identical

b151b91

add client_id to JobHeaders

dd532ff

fix formatting

0451b5d

fix failing test

1f497fb

address cl comments

e5953da

fix tests failing after rebase

7e7eb1f

fix merge errors, use the new vendored deps

6e87f84

address cl comments

df8148c

m-trieu added 6 commits February 13, 2024 10:14

override authority of grpc channel for direct path

a348d9d

Do not inject WindmillServer from deps, create in StreamingDataflowWo…

c9d5cc0

…rker

fix failing test

44c33a3

rebase, factor out FailedTokens class

e2e75a4

address cl comments

18ee453

address cl comments

2b59492

m-trieu force-pushed the mt-metadata-api branch from a60ec02 to 2b59492 Compare February 13, 2024 19:26

scwhittle approved these changes Feb 15, 2024

View reviewed changes

address comments

d9cbe0f

scwhittle approved these changes Feb 15, 2024

View reviewed changes

scwhittle merged commit 690a5a4 into apache:master Feb 15, 2024
17 checks passed

		@@ -45,92 +48,156 @@ class GrpcDispatcherClient {
		private final WindmillStubFactory windmillStubFactory;

		@GuardedBy("this")

Update windmill proto definition #30046

Update windmill proto definition #30046

Conversation

m-trieu commented Jan 19, 2024 • edited Loading

r: @scwhittle

GitHub Actions Tests Status (on master branch)

github-actions bot commented Jan 19, 2024

m-trieu commented Jan 19, 2024

m-trieu commented Jan 19, 2024

github-actions bot commented Jan 19, 2024

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

lostluck commented Jan 19, 2024

github-actions bot commented Jan 19, 2024

m-trieu commented Jan 19, 2024

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

m-trieu commented Feb 9, 2024

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

m-trieu commented Feb 15, 2024

scwhittle left a 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

m-trieu commented Jan 19, 2024 •

edited

Loading