Repository navigation
Support Environment Sensitive Staging - #40699
tarun-google wants to merge 8 commits into
Conversation
There was a problem hiding this comment.
📋 Review Summary
This PR fixes artifact staging for Dataflow jobs using prebuilt SDK container images (--prebuild_sdk_container_engine) by decoupling pipeline environment artifact dependencies from the job-level staged resources in apiclient.py, and introduces a fallback in boot.go to copy pre-staged artifacts from /opt/apache/beam/artifacts when no artifacts are provided via the Provision API.
🔍 General Feedback
- The separation of
dependencies(for pipeline proto environments) andresources(for Dataflow job staging) inapiclient.pycleanly resolves the regression whereStager.create_job_resourceswas skipped during job submission. - However, the fallback in
boot.gohas a couple of edge cases regardingsubmission_environment_dependencies.txt(which is always staged byapiclient.pyeven with prebuilt containers) and potential filename collisions across multiple prebuilt image layers.
💡 Tip: Comment /gemini review on this pull request to trigger a new Gemini code review.
💬 Inline Findings
sdks/python/container/boot.go(line 270): 🟠 High:len(info.GetDependencies()) == 0will evaluate tofalseon Dataflow even when--prebuild_sdk_container_engineis used, preventing this fallback from ever executing on Dataflow. Inapiclient.py,_stage_resourcescallsStager.create_job_resourceswith the defaultlog_submission_env_dependencies=True, which stagessubmission_environment_dependencies.txtto GCS. Dataflow's runner harness populatesProvisionInfo.dependenciesfrom the job's GCSpackageslist (not from the pipeline proto'sEnvironment.dependencies), meaninginfo.GetDependencies()will containsubmission_environment_dependencies.txt(length 1) rather than being empty.
Additionally, even if /opt/apache/beam/artifacts is copied to dir, installSetupPackages only iterates over the files slice populated from artifacts (lines 296-304). Any workflow.tar.gz or PICKLED_MAIN_SESSION copied from /opt/apache/beam/artifacts won't be in files and thus won't be processed by installSetupPackages (though PICKLED_MAIN_SESSION in dir is still picked up by sdk_worker_main.py). Consider checking if /opt/apache/beam/artifacts exists and merging/copying any local artifacts not already present in dir.
-
sdks/python/container/boot.go(line 484): 🟡 Medium:copyDirwalkssrcrecursively (which includes subdirectories like/opt/apache/beam/artifacts/<hash>/...created byboot --setup_onlyduring the image build), but flattens all files into the top-leveldstdirectory usingfilepath.Base(path). Becauseboot --setup_onlyalready copies artifacts from/opt/apache/beam/artifacts/<name>into/opt/apache/beam/artifacts/<hash>/<name>, every artifact exists at least twice under/opt/apache/beam/artifacts. If a user builds a multi-layered prebuilt image, older subdirectories could overwrite newer top-level files depending on lexical walk order, or if any subdirectory exists,os.Open(path)followed byio.Copyon a directory entry will fail or produce unexpected behavior ifinfo.IsDir()wasn't skipped (thoughinfo.IsDir()is skipped, nested files in subdirectories will still collide onfilepath.Base(path)). Consider usingos.ReadDir(src)instead offilepath.Walkto only copy regular files in the top-level/opt/apache/beam/artifactsdirectory (and skipartifacts_info.json). -
sdks/python/container/boot.go(line 500): 🟢 Low:defer out.Close()ignores any write/flush errors that occur when closing the destination file, and keepingdeferinside thefilepath.Walkcallback holds all file descriptors open untilcopyDirreturns. Closeinandoutexplicitly per iteration and check the error fromout.Close().
Codecov Report❌ Patch coverage is
Additional details and impacted files@@ Coverage Diff @@
## master #40699 +/- ##
===========================================
- Coverage 80.04% 56.42% -23.62%
===========================================
Files 490 875 +385
Lines 82916 154692 +71776
===========================================
+ Hits 66373 87291 +20918
- Misses 16543 65545 +49002
- Partials 0 1856 +1856
Flags with carried forward coverage won't be shown. Click here to find out more. ☔ View full report in Codecov by Harness. 🚀 New features to boost your workflow:
|
|
Assigning reviewers: R: @claudevdm for label python. 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). |
claudevdm
left a comment
There was a problem hiding this comment.
Can you trigger post commits?
|
R: @tvalentyn can you PTAL too since this has a large blast radius. |
|
Stopping reviewer notifications for this pull request: review requested by someone other than the bot, ceding control. If you'd like to restart, comment |
|
Run Python PostCommit |
|
I do not think above command triggered Post-Submits. What is the process? Also, we should get this before next release cut. cc: @tvalentyn |
|
@tarun-google you have to make any trivial modification to a trigger file corresponding to the postcommit suite, and add it to the pr. I haven't done a deep dive on this PR yet, but it's on my todo list. |
|
|
||
| os.Setenv("SEMI_PERSISTENT_DIRECTORY", *semiPersistDir) | ||
| // Point the SDK worker at the per-environment root so that it keeps | ||
| // resolving <SEMI_PERSISTENT_DIRECTORY>/staged/... (e.g. pickled_main_session). |
There was a problem hiding this comment.
My main concern is: some users might have made assumptions where files are staged, now the path becomes dynamic, so we need to evaluate it on the worker and not at pipeline construction. We need to communicate that this is a breaking change for Python in CHANGES.md and guide users how to navigate this change.
I think any code that hardcodes /tmp/staged/... (non-dataflow) or smth like /var/opt/google/staged/... will fail to find files in the right location.
We can point users to a code snippet https://beam.apache.org/documentation/sdks/python-pipeline-dependencies/#accessing-staged-files-on-the-workers (sent you a PR correcting it slightly).
We might need to highlight that this evaluation needs to run on the workers, and not at pipeline submission.
We should check how would users use --files_to_stage to supply a model, and then point RunInference transform to the model that was staged?
We can start with checking if these tests will fail
There was a problem hiding this comment.
Yes, if some already know that main requirements are replacing the provider and did not supply all the required dependencies for transforms and included then in main requirements to form super set. The pipeline would break with no module error.
There was a problem hiding this comment.
Done, Updated CHANGES.md
There was a problem hiding this comment.
hmm let me check the tests. I do not think users should know and care about how workers are downloading and managing dependencies. Tests We have written, I agree we can find a way to fix them.
There was a problem hiding this comment.
@tvalentyn Created #40795 for fixing the tests
DOC: https://docs.google.com/document/d/1gNfjbIYNzbv15-RjCj0ad4MHenVPoSjUEcH0Ok3iNyU/edit?usp=sharing
Tested dependencies are downloaded to separate directories by sdk container

requirements.txt are appended with sha-256 to avoid replacing
