From cb8eed8535a9b02b1a1cd958fe2d4feccc607d71 Mon Sep 17 00:00:00 2001 From: Reuven Lax Date: Mon, 28 Sep 2026 11:21:31 -0700 Subject: [PATCH 1/3] Share and cache GCS staging across TestDataflowRunner executions --- .../google-cloud-dataflow-java/build.gradle | 2 ++ .../runners/dataflow/TestDataflowRunner.java | 6 ++++++ .../beam/runners/dataflow/util/GcsStager.java | 17 +++++++++++++++++ .../dataflow/TestDataflowRunnerTest.java | 18 ++++++++++++++++++ 4 files changed, 43 insertions(+) diff --git a/runners/google-cloud-dataflow-java/build.gradle b/runners/google-cloud-dataflow-java/build.gradle index fbfa837d9fd6..aede28ded835 100644 --- a/runners/google-cloud-dataflow-java/build.gradle +++ b/runners/google-cloud-dataflow-java/build.gradle @@ -173,6 +173,7 @@ def legacyPipelineOptions = [ "--project=${gcpProject}", "--region=${gcpRegion}", "--tempRoot=${dataflowValidatesTempRoot}", + "--stagingLocation=${dataflowValidatesTempRoot}/staging", "--dataflowWorkerJar=${dataflowLegacyWorkerJar}", "--usePublicIps=false", "--experiments=enable_lineage" @@ -190,6 +191,7 @@ def runnerV2CommonPipelineOptions = [ "--project=${gcpProject}", "--region=${gcpRegion}", "--tempRoot=${dataflowValidatesTempRoot}", + "--stagingLocation=${dataflowValidatesTempRoot}/staging", "--experiments=use_unified_worker,use_runner_v2", "--firestoreDb=${firestoreDb}", "--usePublicIps=false", diff --git a/runners/google-cloud-dataflow-java/src/main/java/org/apache/beam/runners/dataflow/TestDataflowRunner.java b/runners/google-cloud-dataflow-java/src/main/java/org/apache/beam/runners/dataflow/TestDataflowRunner.java index db8364bcbe8e..89d198154ad3 100644 --- a/runners/google-cloud-dataflow-java/src/main/java/org/apache/beam/runners/dataflow/TestDataflowRunner.java +++ b/runners/google-cloud-dataflow-java/src/main/java/org/apache/beam/runners/dataflow/TestDataflowRunner.java @@ -85,6 +85,12 @@ public static TestDataflowRunner fromOptions(PipelineOptions options) { tempLocation = tempLocation.substring(0, tempLocation.length() - File.separator.length()); } dataflowOptions.setTempLocation(tempLocation); + String defaultPerJobStagingLocation = + FileSystems.matchNewDirectory(tempLocation, "staging").toString(); + if (defaultPerJobStagingLocation.equals(dataflowOptions.getStagingLocation())) { + dataflowOptions.setStagingLocation( + FileSystems.matchNewDirectory(dataflowOptions.getTempRoot(), "staging").toString()); + } return new TestDataflowRunner( dataflowOptions, DataflowClient.create(options.as(DataflowPipelineOptions.class))); diff --git a/runners/google-cloud-dataflow-java/src/main/java/org/apache/beam/runners/dataflow/util/GcsStager.java b/runners/google-cloud-dataflow-java/src/main/java/org/apache/beam/runners/dataflow/util/GcsStager.java index bf34e007c40c..a5cda3274b18 100644 --- a/runners/google-cloud-dataflow-java/src/main/java/org/apache/beam/runners/dataflow/util/GcsStager.java +++ b/runners/google-cloud-dataflow-java/src/main/java/org/apache/beam/runners/dataflow/util/GcsStager.java @@ -21,7 +21,11 @@ import static org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.Preconditions.checkArgument; import com.google.api.services.dataflow.model.DataflowPackage; +import java.util.Arrays; +import java.util.Collections; import java.util.List; +import java.util.concurrent.ConcurrentHashMap; +import org.apache.beam.runners.dataflow.TestDataflowRunner; import org.apache.beam.runners.dataflow.options.DataflowPipelineOptions; import org.apache.beam.runners.dataflow.util.PackageUtil.StagedFile; import org.apache.beam.sdk.extensions.gcp.storage.GcsCreateOptions; @@ -30,6 +34,9 @@ /** Utility class for staging files to GCS. */ public class GcsStager implements Stager { + private static final ConcurrentHashMap, List> STAGED_FILES_CACHE = + new ConcurrentHashMap<>(); + private DataflowPipelineOptions options; private GcsStager(DataflowPipelineOptions options) { @@ -49,6 +56,16 @@ public static GcsStager fromOptions(PipelineOptions options) { */ @Override public List stageFiles(List filesToStage) { + String stagingLocation = options.getStagingLocation(); + if (stagingLocation != null && TestDataflowRunner.class.equals(options.getRunner())) { + List cacheKey = Arrays.asList(stagingLocation, filesToStage); + return STAGED_FILES_CACHE.computeIfAbsent( + cacheKey, k -> Collections.unmodifiableList(stageFilesUncached(filesToStage))); + } + return stageFilesUncached(filesToStage); + } + + private List stageFilesUncached(List filesToStage) { try (PackageUtil packageUtil = PackageUtil.withDefaultThreadPool()) { return packageUtil.stageClasspathElements( filesToStage, options.getStagingLocation(), buildCreateOptions()); diff --git a/runners/google-cloud-dataflow-java/src/test/java/org/apache/beam/runners/dataflow/TestDataflowRunnerTest.java b/runners/google-cloud-dataflow-java/src/test/java/org/apache/beam/runners/dataflow/TestDataflowRunnerTest.java index ed6259a3ee22..3bac63fecca8 100644 --- a/runners/google-cloud-dataflow-java/src/test/java/org/apache/beam/runners/dataflow/TestDataflowRunnerTest.java +++ b/runners/google-cloud-dataflow-java/src/test/java/org/apache/beam/runners/dataflow/TestDataflowRunnerTest.java @@ -102,6 +102,24 @@ public void testToString() { "TestDataflowRunner#TestAppName", TestDataflowRunner.fromOptions(options).toString()); } + @Test + public void testFromOptionsUsesSharedStagingLocationUnderTempRoot() { + options.setJobName("test-job-1"); + TestDataflowRunner.fromOptions(options); + assertEquals("gs://test/test-job-1/output/results", options.getTempLocation()); + assertEquals("gs://test/test-job-1/output/results", options.getGcpTempLocation()); + assertEquals("gs://test/staging/", options.getStagingLocation()); + } + + @Test + public void testFromOptionsPreservesExplicitStagingLocation() { + options.setJobName("test-job-2"); + options.setStagingLocation("gs://custom-bucket/custom-staging/"); + TestDataflowRunner.fromOptions(options); + assertEquals("gs://test/test-job-2/output/results", options.getTempLocation()); + assertEquals("gs://custom-bucket/custom-staging/", options.getStagingLocation()); + } + @Test public void testRunBatchJobThatSucceeds() throws Exception { Pipeline p = Pipeline.create(options); From d8b225a46fc21d63500a684b93fce432eca6e67a Mon Sep 17 00:00:00 2001 From: Reuven Lax Date: Mon, 28 Sep 2026 13:13:30 -0700 Subject: [PATCH 2/3] fixes --- .../beam/runners/dataflow/util/GcsStager.java | 19 +++++++++++++++---- .../dataflow/TestDataflowRunnerTest.java | 2 ++ 2 files changed, 17 insertions(+), 4 deletions(-) diff --git a/runners/google-cloud-dataflow-java/src/main/java/org/apache/beam/runners/dataflow/util/GcsStager.java b/runners/google-cloud-dataflow-java/src/main/java/org/apache/beam/runners/dataflow/util/GcsStager.java index a5cda3274b18..20fbbf221b32 100644 --- a/runners/google-cloud-dataflow-java/src/main/java/org/apache/beam/runners/dataflow/util/GcsStager.java +++ b/runners/google-cloud-dataflow-java/src/main/java/org/apache/beam/runners/dataflow/util/GcsStager.java @@ -21,7 +21,7 @@ import static org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.Preconditions.checkArgument; import com.google.api.services.dataflow.model.DataflowPackage; -import java.util.Arrays; +import com.google.auto.value.AutoValue; import java.util.Collections; import java.util.List; import java.util.concurrent.ConcurrentHashMap; @@ -34,8 +34,19 @@ /** Utility class for staging files to GCS. */ public class GcsStager implements Stager { - private static final ConcurrentHashMap, List> STAGED_FILES_CACHE = - new ConcurrentHashMap<>(); + @AutoValue + abstract static class StagedFilesCacheKey { + abstract String getStagingLocation(); + + abstract List getFilesToStage(); + + static StagedFilesCacheKey of(String stagingLocation, List filesToStage) { + return new AutoValue_GcsStager_StagedFilesCacheKey(stagingLocation, filesToStage); + } + } + + private static final ConcurrentHashMap> + STAGED_FILES_CACHE = new ConcurrentHashMap<>(); private DataflowPipelineOptions options; @@ -58,7 +69,7 @@ public static GcsStager fromOptions(PipelineOptions options) { public List stageFiles(List filesToStage) { String stagingLocation = options.getStagingLocation(); if (stagingLocation != null && TestDataflowRunner.class.equals(options.getRunner())) { - List cacheKey = Arrays.asList(stagingLocation, filesToStage); + StagedFilesCacheKey cacheKey = StagedFilesCacheKey.of(stagingLocation, filesToStage); return STAGED_FILES_CACHE.computeIfAbsent( cacheKey, k -> Collections.unmodifiableList(stageFilesUncached(filesToStage))); } diff --git a/runners/google-cloud-dataflow-java/src/test/java/org/apache/beam/runners/dataflow/TestDataflowRunnerTest.java b/runners/google-cloud-dataflow-java/src/test/java/org/apache/beam/runners/dataflow/TestDataflowRunnerTest.java index 3bac63fecca8..4f6ff01c327a 100644 --- a/runners/google-cloud-dataflow-java/src/test/java/org/apache/beam/runners/dataflow/TestDataflowRunnerTest.java +++ b/runners/google-cloud-dataflow-java/src/test/java/org/apache/beam/runners/dataflow/TestDataflowRunnerTest.java @@ -49,6 +49,7 @@ import org.apache.beam.sdk.extensions.gcp.auth.TestCredential; import org.apache.beam.sdk.extensions.gcp.storage.NoopPathValidator; import org.apache.beam.sdk.extensions.gcp.util.Transport; +import org.apache.beam.sdk.io.FileSystems; import org.apache.beam.sdk.options.PipelineOptionsFactory; import org.apache.beam.sdk.testing.PAssert; import org.apache.beam.sdk.testing.SerializableMatcher; @@ -94,6 +95,7 @@ public void setUp() throws Exception { options.setGcpCredential(new TestCredential()); options.setRunner(TestDataflowRunner.class); options.setPathValidatorClass(NoopPathValidator.class); + FileSystems.setDefaultPipelineOptions(options); } @Test From d2941d54419fb4929bec6f904b538bde5189918e Mon Sep 17 00:00:00 2001 From: Reuven Lax Date: Tue, 29 Sep 2026 08:38:22 -0700 Subject: [PATCH 3/3] address comments --- .../beam/runners/dataflow/util/GcsStager.java | 31 ++++++++++++++----- 1 file changed, 24 insertions(+), 7 deletions(-) diff --git a/runners/google-cloud-dataflow-java/src/main/java/org/apache/beam/runners/dataflow/util/GcsStager.java b/runners/google-cloud-dataflow-java/src/main/java/org/apache/beam/runners/dataflow/util/GcsStager.java index 20fbbf221b32..413a870a365b 100644 --- a/runners/google-cloud-dataflow-java/src/main/java/org/apache/beam/runners/dataflow/util/GcsStager.java +++ b/runners/google-cloud-dataflow-java/src/main/java/org/apache/beam/runners/dataflow/util/GcsStager.java @@ -22,15 +22,19 @@ import com.google.api.services.dataflow.model.DataflowPackage; import com.google.auto.value.AutoValue; +import java.time.Duration; import java.util.Collections; import java.util.List; -import java.util.concurrent.ConcurrentHashMap; -import org.apache.beam.runners.dataflow.TestDataflowRunner; +import java.util.concurrent.ExecutionException; import org.apache.beam.runners.dataflow.options.DataflowPipelineOptions; import org.apache.beam.runners.dataflow.util.PackageUtil.StagedFile; import org.apache.beam.sdk.extensions.gcp.storage.GcsCreateOptions; import org.apache.beam.sdk.options.PipelineOptions; import org.apache.beam.sdk.util.MimeTypes; +import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.Throwables; +import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.cache.Cache; +import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.cache.CacheBuilder; +import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.util.concurrent.UncheckedExecutionException; /** Utility class for staging files to GCS. */ public class GcsStager implements Stager { @@ -45,8 +49,13 @@ static StagedFilesCacheKey of(String stagingLocation, List filesToSt } } - private static final ConcurrentHashMap> - STAGED_FILES_CACHE = new ConcurrentHashMap<>(); + private static final int MAX_STAGED_FILES_CACHE_SIZE = 5000; + + private static final Cache> STAGED_FILES_CACHE = + CacheBuilder.newBuilder() + .maximumSize(MAX_STAGED_FILES_CACHE_SIZE) + .expireAfterWrite(Duration.ofMinutes(30)) + .build(); private DataflowPipelineOptions options; @@ -68,10 +77,18 @@ public static GcsStager fromOptions(PipelineOptions options) { @Override public List stageFiles(List filesToStage) { String stagingLocation = options.getStagingLocation(); - if (stagingLocation != null && TestDataflowRunner.class.equals(options.getRunner())) { + if (stagingLocation != null) { StagedFilesCacheKey cacheKey = StagedFilesCacheKey.of(stagingLocation, filesToStage); - return STAGED_FILES_CACHE.computeIfAbsent( - cacheKey, k -> Collections.unmodifiableList(stageFilesUncached(filesToStage))); + try { + return STAGED_FILES_CACHE.get( + cacheKey, () -> Collections.unmodifiableList(stageFilesUncached(filesToStage))); + } catch (ExecutionException | UncheckedExecutionException e) { + if (e.getCause() != null) { + Throwables.throwIfUnchecked(e.getCause()); + throw new RuntimeException(e.getCause()); + } + throw new RuntimeException(e); + } } return stageFilesUncached(filesToStage); }