add computation state cache, config loaders, and status pages #31133

m-trieu · 2024-04-30T11:40:57Z

Add computation state cache and config loaders that load the computation state cache
move status pages out of StreamingDataflowWorker file

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-04-30T12:07:29Z

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

scwhittle

Some initial comments, didn't get through everything

scwhittle · 2024-04-30T13:16:53Z

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

            sampler,
            metricTrackingWindmillServer::refreshActiveWork,
            executorSupplier.apply("RefreshWork"));
+
+    WorkerStatusPages workerStatusPages =
+        WorkerStatusPages.create(DEFAULT_STATUS_PORT, memoryMonitor, () -> true);


can this last parameter be made optional via overload? or at least add a comment here

done. overloaded

scwhittle · 2024-04-30T13:19:13Z

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

+        WorkerStatusPages.create(DEFAULT_STATUS_PORT, memoryMonitor, () -> true);
+    this.statusPages =
+        windmillServiceEnabled
+            ? StreamingWorkerStatusPages.forStreamingEngine(


would a builder pattern be better than the separate forStreamingEngine and forAppliance factory methods?

There seem to be a lot of common things.

scwhittle · 2024-04-30T13:20:17Z

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

+    WindmillServerStub windmillServer = createWindmillServerStub(options, windmillStreamFactory);
+    ComputationConfig.Fetcher configFetcher =
+        options.isEnableStreamingEngine()
+            ? StreamingEngineConfigFetcher.forTesting(


forTesting isn't a good name if this is in the regular path

removed this was supposed to just use the create method

scwhittle · 2024-04-30T13:20:44Z

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

+    ComputationConfig.Fetcher configFetcher =
+        options.isEnableStreamingEngine()
+            ? StreamingEngineConfigFetcher.forTesting(
+                true,


rm param if hard-coded

it is used in a different test

scwhittle · 2024-04-30T13:32:10Z

...n/java/org/apache/beam/runners/dataflow/worker/streaming/config/StreamingPipelineConfig.java

+        .setWindmillServiceEndpoints(ImmutableSet.of());
+  }
+
+  public static StreamingPipelineConfig forAppliance(


can we just get rid of this and have the appliance config use builder() and set these?

scwhittle · 2024-04-30T13:43:37Z

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

+                        maxWorkItemCommitBytes))
+            : new StreamingApplianceConfigFetcher(
+                windmillServer,
+                config -> consumeUserStepToStateFamilyName(config, stateNameMap),


seems like the consuming the config method could be the same for both SE and appliance, where unset fields are just ignored

scwhittle · 2024-04-30T13:44:53Z

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

+    ComputationStateCache computationStateCache =
+        ComputationStateCache.create(
+            configFetcher, workExecutor, windmillStateCache::forComputation);
+    if (windmillServer instanceof GrpcWindmillServer) {


this cast and lazy initialization is kind of gross

one idea could be to remove the windmillserver from the appliance config fetcher, it could just construct it's own channel and sync stub directly.
and then have a separate start method on the configfetcher taking the function to consume responses. Then you can make sure to call that after everything is initialized.

sgtm agreed

i think the issue would be if it is using the WindmillServerBase implementation of WindmillServerStub, which uses JNI

i will try to find a way around this

i actually realized that appliance does not refresh heartbeats which is what that consumer of heartbeat responses is for.

we only call refresh active work on get data stream if streaming engine is enabled.

done
i changed the way this was created, still not super clean but much better than before.

I wonder if we can make GrpcWindmillServer not implement WindmillServerStub. Not sure it currently makes sense to group together Appliance and Engine client implementations and leads to wonky situtations like above.

m-trieu · 2024-05-01T13:56:58Z

failing test is unrelated and passes on local runs

scwhittle · 2024-05-02T13:02:19Z

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

+            .collect(Collectors.toList()),
+        workExecutor,
+        stateCache::forComputation);
+    computationStateCacheRef.set(computationStateCache);


how about a VisibleForTesting annotated method to access the ComputationStateCache instead? Seems convoluted to pass in atomic ref to assign to.

scwhittle · 2024-05-02T13:09:19Z

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

-          60,
-          TimeUnit.SECONDS);
-      scheduledExecutors.add(statusPageTimer);
-    }
    workCommitter.start();
    workerStatusReporter.start();
    activeWorkRefresher.start();
  }

  public void startStatusPages() {


can this be inlined into start() and removed?

today we don't start them for tests (same across batch and streaming harnesses)

Add a comment that this may be omitted for lighterweight testing?

scwhittle · 2024-05-02T13:12:29Z

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

+                  Collection<LatencyAttribution> getWorkStreamLatencies) ->
+                  computationStateCache
+                      .get(computation)
+                      .ifPresent(


what happened before on missing computation? In either case it seems like we should throw exception or log error as otherwise we're just dropping the work item silently which will be confusing.

added more logging in the ComputationStateCache

public Optional<ComputationState> get(String computationId) { try { return Optional.ofNullable(computationCache.get(computationId)); } catch (ExecutionException | ComputationStateNotFoundException e) { if (e.getCause() instanceof ComputationStateNotFoundException) { LOG.error( "Trying to fetch unknown computation={}, known computations are {}.", computationId, getAllComputationIds()); } else { LOG.warn("Error occurred fetching computation for computationId={}", computationId, e); } }

scwhittle · 2024-05-02T13:21:10Z

...rg/apache/beam/runners/dataflow/worker/streaming/config/StreamingApplianceConfigFetcher.java

+/** Fetches computation config from Streaming Appliance. */
+@Internal
+@ThreadSafe
+public final class StreamingApplianceConfigFetcher implements ComputationConfig.Fetcher {


should this be StreamingApplianceComputationConfigFetcher? ditto with SE

scwhittle · 2024-05-02T13:24:55Z

...rg/apache/beam/runners/dataflow/worker/streaming/config/StreamingApplianceConfigFetcher.java

+      Consumer<StreamingPipelineConfig> onPipelineConfig,
+      Function<MapTask, MapTask> fixMapTaskMultiOutputInfoFn) {
+    this.windmillServer = windmillServer;
+    this.onPipelineConfig = onPipelineConfig;


seems odd to have this listener hidden in the fetcher

Maybe instead the listener should be on whatever is driving the fetching. I think that woudl be the cache.

Additionally I think that the FIX_MULTI_OUTPUT_INFOS_ON_PAR_DO_INSTRUCTIONS could move into the cache where fetches are being performed. As is it is injected into all teh fetchers and applied to the testing generated configs. That would let you just do it once and possibly coudl remove from StreamingDataflowWorker to the cache

done and moved

changed it to be a member instead of static since we need the static global id generator from StreamingDataflowWorker.

scwhittle · 2024-05-06T13:02:43Z

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

-          60,
-          TimeUnit.SECONDS);
-      scheduledExecutors.add(statusPageTimer);
-    }
    workCommitter.start();
    workerStatusReporter.start();
    activeWorkRefresher.start();
  }

  public void startStatusPages() {


Add a comment that this may be omitted for lighterweight testing?

scwhittle · 2024-05-06T13:03:07Z

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

-                !computationState.getTransformUserNameToStateFamily().isEmpty()
-                    ? computationState.getTransformUserNameToStateFamily()
-                    : stateNameMap,
+                //                !computationState.getTransformUserNameToStateFamily().isEmpty()


rm comments

scwhittle · 2024-05-06T13:16:18Z

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

+      ConcurrentMap<String, String> stateNameMap) {
+    Function<MapTask, MapTask> fixMultiOutputInfosOnParDoInstructions =
+        new FixMultiOutputInfosOnParDoInstructions(idGenerator);
+    return new ComputationStateCache(


Could this instead just use create and then poke in to add initial state name map values?
via
cache = create(...);
cache.stateNameMap.addAll(stateNameMap);

It's a fair bit of setup to duplicate and it coudl drift meaning we're not testing the main code-path.

Or you could further reduce duplication of the parameters the create/forTesting methods and just expose a visiblefortesting method for the statenamemap and the current caller of this method could use that directly.

scwhittle · 2024-05-06T13:22:10Z

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

+  }
+
+  /** Returns a read-only view of all computations. */
+  public ImmutableList<ComputationState> getAllComputations() {


should this be getAllPresentComputations?

scwhittle · 2024-05-06T13:26:27Z

...eam/runners/dataflow/worker/streaming/config/StreamingApplianceComputationConfigFetcher.java

+        "computationId is empty. Cannot fetch computation config without a computationId.");
+
+    GetConfigResponse response =
+        windmillServer.getConfig(


instead of taking in WindmillServer can you take in a functional interface matching this method?

That will help show that the full class isn't used and could simplify testing. It could help in the future break up windmillServer if we get rid of the jni class keeping it all tied together ATM.

scwhittle · 2024-05-06T13:40:27Z

...rc/main/java/org/apache/beam/runners/dataflow/worker/streaming/config/ComputationConfig.java

+
+    default void stop() {}
+
+    Optional<ComputationConfig> getConfig(String computationId);


let's call this fetchConfig, it makes it clearer that it is likely an rpc (and matches class name)

scwhittle · 2024-05-06T13:42:49Z

...e/beam/runners/dataflow/worker/streaming/config/StreamingEngineComputationConfigFetcher.java

+  private Optional<StreamingEnginePipelineConfig> getComputationConfigInternal(
+      String computationId) {
+    Optional<StreamingEnginePipelineConfig> streamingConfig = getConfigInternal(computationId);
+    streamingConfig.ifPresent(onStreamingConfig);


can we remove this listening from the computation config fetching? It seems like it should just be the global config from the background thread that runs onStreamingConfig.

scwhittle · 2024-05-06T13:47:35Z

...e/beam/runners/dataflow/worker/streaming/config/StreamingEngineComputationConfigFetcher.java

+    LOG.info("Initial global configuration received, harness is now ready");
+  }
+
+  private Optional<StreamingEnginePipelineConfig> getComputationConfigInternal(


I think this should just return the computationconfig not full StreamingEnginePipelineConfig. I think that StreamingEnginePipelineConfig can have the computation config removed from it, as we don't care about listening to computations as they are fetched.

That will then make it clearer that we only want the listener to trigger on the periodic background fetching, not the computation fetching driven by the cache.

scwhittle · 2024-05-06T13:49:38Z

.../org/apache/beam/runners/dataflow/worker/streaming/config/StreamingEnginePipelineConfig.java

+
+  public abstract Map<String, String> userStepToStateFamilyNameMap();
+
+  public abstract Optional<StreamingComputationConfig> computationConfig();


see other comment, I think this class should just be for the global config, not related to a computation.

scwhittle · 2024-05-06T13:51:32Z

...e/beam/runners/dataflow/worker/streaming/config/StreamingEngineComputationConfigFetcher.java

+    return getConfigInternal(null);
+  }
+
+  private Optional<StreamingEnginePipelineConfig> getConfigInternal(@Nullable String computation) {


probably need to change this to some template if we're returnign different types for global or per-computation config

github-actions · 2024-05-07T06:36:09Z

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

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

scwhittle

Looking pretty good, mostly just some test comments

scwhittle · 2024-05-07T08:49:09Z

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

@@ -54,6 +56,7 @@ public final class ComputationStateCache implements StatusDataProvider {

  private static final Logger LOG = LoggerFactory.getLogger(ComputationStateCache.class);

+  private final ConcurrentMap<String, String> globalUsernameToStateFamilyNameMap;


nit: global has connotations and makes it sound like static global, how about pipelineUserNameToStateFamilyNameMap

scwhittle · 2024-05-07T08:52:57Z

...e/beam/runners/dataflow/worker/streaming/config/StreamingEngineComputationConfigFetcher.java

+    }
+  }
+
+  private static Optional<StreamingEnginePipelineConfig> createPipelineConfig(


don't think this needs to return optional

scwhittle · 2024-05-07T09:22:54Z

...orker/src/test/java/org/apache/beam/runners/dataflow/worker/StreamingDataflowWorkerTest.java

+        localRetryTimeoutMs);
+  }
+
+  private StreamingDataflowWorker makeWorker(


there are 6 makeWorkers, it seems some autovalue builder in this test would make it more readable and less duplication of the default for various params.

done used a builder with some test parameters

scwhittle · 2024-05-07T09:25:24Z

...c/test/java/org/apache/beam/runners/dataflow/worker/streaming/ComputationStateCacheTest.java

+  private ComputationStateCache computationStateCache;
+
+  private static Work createWork(long workToken, long cacheToken) {
+


nit: rm blank line

...c/test/java/org/apache/beam/runners/dataflow/worker/streaming/ComputationStateCacheTest.java

...am/runners/dataflow/worker/streaming/config/StreamingEngineComputationConfigFetcherTest.java

scwhittle · 2024-05-07T09:47:04Z

...am/runners/dataflow/worker/streaming/config/StreamingEngineComputationConfigFetcherTest.java

+    asyncStartConfigLoader.start();
+    numExpectedRefreshes.await();
+    asyncStartConfigLoader.join();
+    assertThat(receivedPipelineConfig)


seems safer to stop the fetcher's background thread before accessing received config (or use a concurrent set)

scwhittle · 2024-05-07T09:47:59Z

...am/runners/dataflow/worker/streaming/config/StreamingEngineComputationConfigFetcherTest.java

+    when(mockDataflowServiceClient.getGlobalStreamingConfigWorkItem())
+        .thenReturn(Optional.of(firstConfig))
+        .thenReturn(Optional.of(secondConfig))
+        .thenReturn(Optional.of(thirdConfig));


will there be a flaky error if a 4th fetch comes in before the background thread stops? shoudl some sort of default response be setup?

...am/runners/dataflow/worker/streaming/config/StreamingEngineComputationConfigFetcherTest.java

…tructor StreamingDataflowWorker instance

scwhittle

Just final couple nits. I'm also running some dataflow specific tests to verify

scwhittle · 2024-05-13T08:42:00Z

...orker/src/test/java/org/apache/beam/runners/dataflow/worker/StreamingDataflowWorkerTest.java

+            defaultWorkerParams()
+                .setInstructions(instructions)
+                .publishCounters()
+                .setOptions(createTestingPipelineOptions())


can setOptions(createTestingPipelineOptions()) since it's default

scwhittle · 2024-05-13T08:43:26Z

...orker/src/test/java/org/apache/beam/runners/dataflow/worker/StreamingDataflowWorkerTest.java

@@ -3203,12 +3145,21 @@ public void testExceptionInvalidatesCache() throws Exception {
                WindowingStrategy.globalDefault()),
            makeSinkInstruction(StringUtf8Coder.of(), 1, GlobalWindow.Coder.INSTANCE));

+    defaultWorkerParams()


rm, not used

scwhittle · 2024-05-13T08:44:18Z

...orker/src/test/java/org/apache/beam/runners/dataflow/worker/StreamingDataflowWorkerTest.java

-    StreamingDataflowWorker worker = makeWorker(instructions, options, true /* publishCounters */);
+    StreamingDataflowWorker worker =
+        makeWorker(
+            defaultWorkerParams()


does passing in "--activeWorkRefreshPeriodMillis=100" let you get rid of options?

scwhittle · 2024-05-13T08:44:27Z

...orker/src/test/java/org/apache/beam/runners/dataflow/worker/StreamingDataflowWorkerTest.java

-    StreamingDataflowWorker worker = makeWorker(instructions, options, true /* publishCounters */);
+    StreamingDataflowWorker worker =
+        makeWorker(
+            defaultWorkerParams()


ditto, here and below

scwhittle · 2024-05-13T08:52:35Z

...c/test/java/org/apache/beam/runners/dataflow/worker/streaming/ComputationStateCacheTest.java

    computationState.activateWork(shardedKey, work1);
    computationState.activateWork(shardedKey, work2);
    computationState.activateWork(shardedKey, work3);

+    // Activate 3 Work(s) for computationId2
+    Optional<ComputationState> maybeComputationState2 = computationStateCache.get(computationId);


nit: how about just activating work1 on computationState and work2, work3 on comptuationState2. Then we're not deviating from normal usage of work (shouldn't be activated on multiple/different comptuatoins) in case it causes issues down the line.

github-actions bot added runners dataflow labels Apr 30, 2024

add computation state cache, config loaders, and status pages

b497948

m-trieu force-pushed the mt-computation-state-cache branch from 28a88b4 to b497948 Compare April 30, 2024 11:42

scwhittle requested changes Apr 30, 2024

View reviewed changes

address cl comments, add tests

17ef3db

scwhittle requested changes May 2, 2024

View reviewed changes

m-trieu force-pushed the mt-computation-state-cache branch from 3581abd to f30767c Compare May 3, 2024 17:38

address cl comments

f7f2481

m-trieu force-pushed the mt-computation-state-cache branch from f30767c to f7f2481 Compare May 4, 2024 05:52

scwhittle requested changes May 6, 2024

View reviewed changes

address cl comments

6b2fff8

github-actions bot added the Next Action: Reviewers label May 7, 2024

scwhittle requested changes May 7, 2024

View reviewed changes

address pr comments

e205e91

m-trieu force-pushed the mt-computation-state-cache branch from b670812 to 099081c Compare May 9, 2024 20:24

clean up StreamingDataflowWorkerTest to use builder pattern when cons…

f32844b

…tructor StreamingDataflowWorker instance

m-trieu force-pushed the mt-computation-state-cache branch from 099081c to f32844b Compare May 9, 2024 20:55

scwhittle requested changes May 13, 2024

View reviewed changes

address cl comments to fix tests

2049e0a

scwhittle approved these changes May 14, 2024

View reviewed changes

scwhittle merged commit 31e81ff into apache:master May 14, 2024
17 checks passed


		default void stop() {}

		Optional<ComputationConfig> getConfig(String computationId);


		public abstract Map<String, String> userStepToStateFamilyNameMap();

		public abstract Optional<StreamingComputationConfig> computationConfig();

		@@ -54,6 +56,7 @@ public final class ComputationStateCache implements StatusDataProvider {

		private static final Logger LOG = LoggerFactory.getLogger(ComputationStateCache.class);

		private final ConcurrentMap<String, String> globalUsernameToStateFamilyNameMap;

		private ComputationStateCache computationStateCache;

		private static Work createWork(long workToken, long cacheToken) {

add computation state cache, config loaders, and status pages #31133

add computation state cache, config loaders, and status pages #31133

Conversation

m-trieu commented Apr 30, 2024

R: @scwhittle

GitHub Actions Tests Status (on master branch)

github-actions bot commented Apr 30, 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

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 May 1, 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

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

github-actions bot commented May 7, 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

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

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