Skip to content

Support Environment Sensitive Staging - #40699

Open
tarun-google wants to merge 8 commits into
apache:masterfrom
tarun-google:fix_env_staging
Open

tarun-google wants to merge 8 commits into
apache:masterfrom
tarun-google:fix_env_staging

Conversation

@tarun-google

Copy link
Copy Markdown
Contributor

DOC: https://docs.google.com/document/d/1gNfjbIYNzbv15-RjCj0ad4MHenVPoSjUEcH0Ok3iNyU/edit?usp=sharing

Tested dependencies are downloaded to separate directories by sdk container
image

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

@github-actions github-actions Bot left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

📋 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) and resources (for Dataflow job staging) in apiclient.py cleanly resolves the regression where Stager.create_job_resources was skipped during job submission.
  • However, the fallback in boot.go has a couple of edge cases regarding submission_environment_dependencies.txt (which is always staged by apiclient.py even 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()) == 0 will evaluate to false on Dataflow even when --prebuild_sdk_container_engine is used, preventing this fallback from ever executing on Dataflow. In apiclient.py, _stage_resources calls Stager.create_job_resources with the default log_submission_env_dependencies=True, which stages submission_environment_dependencies.txt to GCS. Dataflow's runner harness populates ProvisionInfo.dependencies from the job's GCS packages list (not from the pipeline proto's Environment.dependencies), meaning info.GetDependencies() will contain submission_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: copyDir walks src recursively (which includes subdirectories like /opt/apache/beam/artifacts/<hash>/... created by boot --setup_only during the image build), but flattens all files into the top-level dst directory using filepath.Base(path). Because boot --setup_only already 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 by io.Copy on a directory entry will fail or produce unexpected behavior if info.IsDir() wasn't skipped (though info.IsDir() is skipped, nested files in subdirectories will still collide on filepath.Base(path)). Consider using os.ReadDir(src) instead of filepath.Walk to only copy regular files in the top-level /opt/apache/beam/artifacts directory (and skip artifacts_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 keeping defer inside the filepath.Walk callback holds all file descriptors open until copyDir returns. Close in and out explicitly per iteration and check the error from out.Close().

@codecov

codecov Bot commented Oct 21, 2026 •

Copy link
Copy Markdown

Codecov Report

❌ Patch coverage is 54.54545% with 5 lines in your changes missing coverage. Please review.
✅ Project coverage is 56.42%. Comparing base (a8d6afd) to head (1870a42).
⚠️ Report is 75 commits behind head on master.

Files with missing lines Patch % Lines
sdks/python/container/boot.go 0.00% 5 Missing ⚠️
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     
Flag Coverage Δ
go 28.90% <0.00%> (?)
python 79.91% <100.00%> (-0.14%) ⬇️

Flags with carried forward coverage won't be shown. Click here to find out more.

☔ View full report in Codecov by Harness.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:
  • ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.
  • 📦 JS Bundle Analysis: Save yourself from yourself by tracking and limiting bundle sizes in JS merges.

@github-actions

Copy link
Copy Markdown
Contributor

Assigning reviewers:

R: @claudevdm for label python.
Backup reviewer: shunping (if the main reviewer is unresponsive, you can escalate to them by tagging their username)

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

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).

@claudevdm claudevdm left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Can you trigger post commits?

@claudevdm

Copy link
Copy Markdown
Collaborator

R: @tvalentyn can you PTAL too since this has a large blast radius.

@github-actions

Copy link
Copy Markdown
Contributor

Stopping reviewer notifications for this pull request: review requested by someone other than the bot, ceding control. If you'd like to restart, comment assign set of reviewers

@tarun-google

Copy link
Copy Markdown
Contributor Author

Run Python PostCommit

@tarun-google

Copy link
Copy Markdown
Contributor Author

I do not think above command triggered Post-Submits. What is the process? Also, we should get this before next release cut. cc: @tvalentyn

@tvalentyn

tvalentyn commented Oct 22, 2026 •

Copy link
Copy Markdown
Contributor

@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.

@tvalentyn

Copy link
Copy Markdown
Contributor

@github-actions github-actions Bot added the size/m M: 50-249 lines changed label Oct 23, 2026

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).

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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

outCol := inference.SklearnModel("/tmp/staged/sklearn_model").RunInference(s, input, inference.WithExpansionAddr(expansionAddr))

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

@tarun-google tarun-google Oct 24, 2026 •

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Done, Updated CHANGES.md

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

@tvalentyn Created #40795 for fixing the tests

Comment thread sdks/python/container/boot.go
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants