Pytorch upgrade in gemma-flex-template docker image
#13262
Conversation
There was a problem hiding this comment.
Hello @svetakvsundhar, I'm Gemini Code Assist1! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed!
Summary of Changes
This pull request upgrades the PyTorch version in the gemma-flex-template Docker image. The base image for the Dockerfile is updated from pytorch/pytorch:2.3.1-cuda11.8-cudnn8-runtime to pytorch/pytorch:2.6.0-cuda11.8-cudnn9-devel. This change resolves a build failure, as evidenced by the screenshots provided in the PR description.
Highlights
- PyTorch Upgrade: The PyTorch version in the Dockerfile is upgraded from 2.3.1 to 2.6.0.
- Base Image Update: The base Docker image is updated to a newer version with CUDA 11.8 and cuDNN 9.
- Build Fix: This change resolves a build failure, as shown in the provided screenshots.
Changelog
- dataflow/gemma-flex-template/Dockerfile
- Updated the base image from
pytorch/pytorch:2.3.1-cuda11.8-cudnn8-runtimetopytorch/pytorch:2.6.0-cuda11.8-cudnn9-develon line 15.
- Updated the base image from
Using Gemini Code Assist
The full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips.
Invoking Gemini
You can request assistance from Gemini at any point in your pull request via creating an issue comment (i.e. comment on the pull request page) using either /gemini <command> or @gemini-code-assist <command>. Below is a summary of the supported commands.
| Feature | Command | Description |
|---|---|---|
| Code Review | /gemini review |
Performs a code review for the current pull request in its current state. |
| Pull Request Summary | /gemini summary |
Provides a summary of the current pull request in its current state. |
| Comment | @gemini-code-assist | Responds in comments when explicitly tagged, both in issue comments and review comments. |
| Help | /gemini help |
Displays a list of available commands. |
Customization
To customize Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a .gemini/ folder in the base of the repository. Detailed instructions can be found here.
Limitations & Feedback
Gemini Code Assist is currently in preview and may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counter productive. You can react with 👍 and 👎 on @gemini-code-assist comments to provide feedback.
You can also get AI-powered code generation, chat, as well as code reviews directly in the IDE at no cost with the Gemini Code Assist IDE Extension.
Did you know?
The name 'PyTorch' is a reference to Torch, a scientific computing framework, with 'Py' signifying its Python binding.
Footnotes
-
Review the Privacy Notices, Generative AI Prohibited Use Policy, Terms of Service, and learn how to configure Gemini Code Assist in GitHub here. Gemini can make mistakes, so double check it and use code with caution. ↩
There was a problem hiding this comment.
Code Review
The pull request upgrades the PyTorch version in the gemma-flex-template Docker image. This change aims to resolve a build failure, as evidenced by the provided screenshots. The upgrade appears to be a straightforward version bump.
Merge Readiness
The pull request appears to address a build failure by upgrading the PyTorch version. Given the successful build screenshot, the change seems effective. However, it's crucial to verify that this upgrade doesn't introduce any regressions or compatibility issues with the existing code and dependencies. Since I am unable to directly approve the pull request, I recommend that others review and approve this code before merging. I recommend merging this pull request after verifying the upgrade doesn't introduce regressions or compatibility issues.
| # This uses Ubuntu with Python 3.10 and comes with CUDA drivers for | ||
| # GPU use. | ||
| ARG SERVING_BUILD_IMAGE=pytorch/pytorch:2.3.1-cuda11.8-cudnn8-runtime | ||
| ARG SERVING_BUILD_IMAGE=pytorch/pytorch:2.6.0-cuda11.8-cudnn9-devel |
There was a problem hiding this comment.
Consider adding a comment explaining why this specific version (2.6.0) was chosen. Was it the latest stable release, or was it selected for compatibility reasons? This context will be helpful for future maintainers. Also, it would be helpful to note the previous version number in the comment for easier comparison.
ARG SERVING_BUILD_IMAGE=pytorch/pytorch:2.6.0-cuda11.8-cudnn9-devel # Upgraded from 2.3.1 to resolve build failure, verified to be compatible
| # This uses Ubuntu with Python 3.10 and comes with CUDA drivers for | ||
| # GPU use. | ||
| ARG SERVING_BUILD_IMAGE=pytorch/pytorch:2.3.1-cuda11.8-cudnn8-runtime | ||
| ARG SERVING_BUILD_IMAGE=pytorch/pytorch:2.6.0-cuda11.8-cudnn9-devel |
There was a problem hiding this comment.
Is there a reason to move from runtime to devel? Are there additional tools the devel image gives that are required here?
There was a problem hiding this comment.
This was the newest Pytorch image, and it worked fine for the build. Would you prefer to stick with runtime images?
There was a problem hiding this comment.
The devel image is more than double the size than runtime, and if there's no additional tools that are required for this sample, it's best to use the simplest image available. The runtime image was also released at the same time as the devel image, only the devel image returns first when sorting alphabetically :)
Co-authored-by: Katie McLaughlin <[email protected]>
Without this change (failing build):
With this change (successful build):
Note: Before submitting a pull request, please open an issue for discussion if you are not associated with Google.
Checklist
nox -s py-3.9(see Test Environment Setup)nox -s lint(see Test Environment Setup)