Repository navigation
Make ValidatesRunner faster: Share and cache GCS staging across TestDataflowRunner executions - #40321
Conversation
| public List<DataflowPackage> stageFiles(List<StagedFile> filesToStage) { | ||
| String stagingLocation = options.getStagingLocation(); | ||
| if (stagingLocation != null && TestDataflowRunner.class.equals(options.getRunner())) { | ||
| List<Object> cacheKey = Arrays.asList(stagingLocation, filesToStage); |
There was a problem hiding this comment.
since the key is a pair of (String, List), consider use it as STAGED_FILES_CACHE's key type. One can use guava's com.google.common.collect.Maps.immutableEntry or org.apache.commons.lang3.tuple.Pair
| "TestDataflowRunner#TestAppName", TestDataflowRunner.fromOptions(options).toString()); | ||
| } | ||
|
|
||
| @Test |
There was a problem hiding this comment.
(Not directly due to this PR) Running the test alone (./gradlew :runners:google-cloud-dataflow-java:test --tests *testFromOptionsUsesSharedStagingLocationUnderTempRoot) it fails with
java.lang.IllegalArgumentException: No filesystem found for scheme gs
at org.apache.beam.sdk.io.FileSystems.getFileSystemInternal(FileSystems.java:558)
at org.apache.beam.sdk.io.FileSystems.matchNewResource(FileSystems.java:665)
at org.apache.beam.sdk.io.FileSystems.matchNewDirectory(FileSystems.java:678)
at org.apache.beam.runners.dataflow.TestDataflowRunner.fromOptions(TestDataflowRunner.java:78)
at org.apache.beam.runners.dataflow.TestDataflowRunnerTest.testFromOptionsUsesSharedStagingLocationUnderTempRoot(TestDataflowRunnerTest.java:108)
but running the test class passed. Likely a pre-existing issue though.
There was a problem hiding this comment.
taking a look....
There was a problem hiding this comment.
it's because when doing this only TestDataflowRunner.fromOptions(options) is called and never Pipeline.create(options). Pipeline.create(options) registers the filesystems as a side effect.
Fixed
| } | ||
|
|
||
| private static final ConcurrentHashMap<StagedFilesCacheKey, List<DataflowPackage>> | ||
| STAGED_FILES_CACHE = new ConcurrentHashMap<>(); |
There was a problem hiding this comment.
Maybe use a Cache with a size bound that will be big enough for any reasonable situation, but will prevent future weirdness if things change.
| @Override | ||
| public List<DataflowPackage> stageFiles(List<StagedFile> filesToStage) { | ||
| String stagingLocation = options.getStagingLocation(); | ||
| if (stagingLocation != null && TestDataflowRunner.class.equals(options.getRunner())) { |
There was a problem hiding this comment.
Why depend on the runner here? Is there a reason it cannot just be always on?
There was a problem hiding this comment.
unlikely to help for other runners (since it's rare to reuse the same JVM for multiple pipelines), however sure - removed this guard.
There was a problem hiding this comment.
fewer deps and lower cyclomatic complexity are also worthwhile :-)
|
Assigning reviewers: R: @tvalentyn added as fallback since no labels match configuration Note: If you would like to opt out of this review, comment Available commands:
The PR bot will only process comments in the main thread (not review comments). |
|
R: @kennknowles |
|
Stopping reviewer notifications for this pull request: review requested by someone other than the bot, ceding control. If you'd like to restart, comment |
TestDataflowRunner previously set tempLocation to //output/results, causing DataflowPipelineOptions.StagingLocationFactory` to derive a unique path for every single test job. Consequently, every test re-hashed ~300+ MB of classpath jars and re-queried/uploaded them to a brand-new GCS staging directory instead of reusing staged artifacts across tests.
Fix: Defaults stagingLocation to a shared /staging/ directory across jobs in the same test run (while keeping tempLocation isolated per job at //output/results).
Caches SHA-256 file hashes and directory zip artifacts in memory keyed by (path, size, lastModifiedMillis) so unmodified classpath jars are hashed once per worker JVM instead of once per test.
Caches resolved DataflowPackage lists per `stagingLocation, filesToStage) when running under TestDataflowRunner so each worker JVM stages/verifies the classpath once.