Use vLLM extra to generate GPU requirements files

[damccorm]

While we can't actually package this in a dockerfile because of licensing issues, this provides a reproducable set of requirements that we can build/test against, and which external users can use to generate their containers.

We can also eventually point to third party providers who publish these containers (e.g. I plan to publish this in a Google-maintained container registry).

Currently, this scopes out Python 3.9 and 3.13

• 3.9 is scoped out because I ran into minor issues with dependency conflicts around transformers versions. These could be fixed, but I didn't bother since we will remove 3.9 support with the next release
• 3.13 requires some wheels to be built on the fly since not all packages provide prebuilt wheels. This is generally fine, but causes issues for some wheels which require rust to be built (this is not part of our CI env, and ideally wouldn't be in most users' containers). I filed [Task]: Add 3.13 GPU containers when dependencies are ready #36637 to track this.

With these changes, I was able to auto-generate #36650 by running the update python dependencies action.

Part of #35487

---

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.

[Amar3tto]

@damccorm I used your changes plus fixed OutOfMemory issue, but there is still an error with starting VLLM server.
Could you please help to understand what happened?
https://github.com/apache/beam/actions/runs/18342904797/job/52242432601

[damccorm]

This PR is unrelated to that flakiness, so I wouldn't expect it to make a difference immediately.

In the logs, I now see a bunch of failures like:

RuntimeError: self and mat2 must have the same dtype, but got Float and Half

I'm not 100% sure what would cause that error, but it may be a consequence of trying to run Gemma on a T4 - could you try updating to an L4 to see if that solves the problem?

[Amar3tto]

This PR is unrelated to that flakiness, so I wouldn't expect it to make a difference immediately.

In the logs, I now see a bunch of failures like:

RuntimeError: self and mat2 must have the same dtype, but got Float and Half

I'm not 100% sure what would cause that error, but it may be a consequence of trying to run Gemma on a T4 - could you try updating to an L4 to see if that solves the problem?

It helped, now I have a green run. Still need to investigate the reasons and perform more tests.
Draft PR: #36451

[codecov]

Codecov Report

✅ All modified and coverable lines are covered by tests.
✅ Project coverage is 55.10%. Comparing base (471050c) to head (3e17728).
⚠️ Report is 15 commits behind head on master.

Additional details and impacted files

@@              Coverage Diff              @@
##             master   #36420       +/-   ##
=============================================
+ Coverage     36.24%   55.10%   +18.85%     
  Complexity     1666     1666               
=============================================
  Files          1060     1059        -1     
  Lines        165704   165727       +23     
  Branches       1195     1195               
=============================================
+ Hits          60063    91320    +31257     
+ Misses       103465    72231    -31234     
  Partials       2176     2176               

Flag	Coverage Δ
python	80.97% <ø> (+40.45%)	⬆️

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

☔ View full report in Codecov by Sentry.
📢 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]

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

[damccorm]

assign set of reviewers

[github-actions]

Assigning reviewers:

R: @tvalentyn for label python.
R: @Abacn for label build.

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

[tvalentyn]

This should be overwritten later

I don't quite follow the comment. Could you please clarify what is being overwritten?

[damccorm]

Thanks - I agree this was unclear. I was trying to say that the torch version we install here doesn't matter since torch will get installed anyways later as part of the vllm install (and that will determine the actual version which is installed). Hopefully the new comment is clearer

[tvalentyn]

Does the outcome of the generated requirements file depend on which wheel is installed (cpu vs non-cpu)? If so, is it because the dependencies of torch change?

[damccorm]

Yes - specifically

beam/sdks/python/container/ml/py310/ml_image_requirements.txt

Line 217 in fe07fe7

torch==2.8.0+cpu

has a cpu postfix, and some gpu requirements (which are not ASF-license compliant) are omitted

[tvalentyn]

Thanks!