Skip to content

Fix Issue 57: Add support for L4 GPUs in Vertex jobs#62

Open
tripathiji1312 wants to merge 1 commit intogoogle-deepmind:mainfrom
tripathiji1312:issue57
Open

Fix Issue 57: Add support for L4 GPUs in Vertex jobs#62
tripathiji1312 wants to merge 1 commit intogoogle-deepmind:mainfrom
tripathiji1312:issue57

Conversation

@tripathiji1312
Copy link
Copy Markdown

This PR resolves issue #57 by adding support for L4 GPUs when launching jobs on Vertex AI.

Description
Currently, attempting to use an L4 GPU by specifying L4_24TH in xm.JobRequirements results in a KeyError: 'NVIDIA_TESLA_L4_24TH'. This is because the accelerator type was being constructed by prefixing the resource name with NVIDIA_TESLA_, which is not the correct identifier for L4 GPUs on Vertex AI. The correct accelerator type is NVIDIA_L4.

This change updates vertex.py to special-case the L4_24TH GPU type and use the correct NVIDIA_L4 accelerator name. A corresponding unit test has been added to vertex_test.py to verify the fix and prevent future regressions.

Changes
Modified vertex.py to correctly handle the L4 GPU accelerator type.
Added test_get_machine_spec_l4 to vertex_test.py to ensure L4 GPUs are correctly configured

@google-cla
Copy link
Copy Markdown

google-cla bot commented Sep 23, 2025

Thanks for your pull request! It looks like this may be your first contribution to a Google open source project. Before we can look at your pull request, you'll need to sign a Contributor License Agreement (CLA).

View this failed invocation of the CLA check for more information.

For the most up to date status, view the checks section at the bottom of the pull request.

@hartikainen
Copy link
Copy Markdown

You might be interested in #58.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants