google-deepmind / google-deepmind/xmanager
Support for L4 GPUs
- Dominant language
- Python
- Stars
- 921
- Forks
- 69
- Avg merge
- 1d 1h
- Merged PRs (30d)
- 1
Description
I see references to `L4` GPUs, e.g. here: https://github.com/google-deepmind/xmanager/blob/c41dc5d15e0a2b18fc1ad6e2d38de59fcc567afd/xmanager/xm/resources.py#L72. This doesn't seem to work though because it's not been implemented. A simple test case in `xmanager/cloud/vertex_test.py` could look something like:
Details
```python
def test_get_machine_spec_l4(self):
job = xm.Job(
executable=local_executables.GoogleContainerRegistryImage('name', ''),
executor=local_executors.Vertex(
requirements=xm.JobRequirements(L4_24TH=2)
),
args={},
)
machine_spec = vertex.get_machine_spec(job)
self.assertDictEqual(
machine_spec,
{
'machine_type': 'g2-standard-4',
'accelerator_type': vertex.aip_v1.AcceleratorType.NVIDIA_L4,
'accelerator_count': 2,
},
)
```
This gives an error:
Details
```console
python -m xmanager.cloud.vertex_test
..........E.Creating CustomJob
CustomJob created. Resource name:
To use this CustomJob in another session:
custom_job = aiplatform.CustomJob.get('')
View Custom Job:
Job launched at:
.
======================================================================
ERROR: test_get_machine_spec_l4 (__main__.VertexTest.test_get_machine_spec_l4)
----------------------------------------------------------------------
Traceback (most recent call last):
File "/Users/user/xmanager/main/xmanager/cloud/vertex_test.py", line 166, in test_get_machine_spec_l4
machine_spec = vertex.get_machine_spec(job)
File "/Users/user/xmanager/main/xmanager/cloud/vertex.py", line 305, in get_machine_spec
spec['accelerator_type'] = aip_v1.AcceleratorType[accelerator_type]
~~~~~~~~~~~~~~~~~~~~~~^^^^^^^^^^^^^^^^^^
File "/opt/homebrew/Cellar/python@3.13/3.13.3/Frameworks/Python.framework/Versions/3.13/lib/python3.13/enum.py", line 791, in __getitem__
return cls._member_map_[name]
~~~~~~~~~~~~~~~~^^^^^^
KeyError: 'NVIDIA_TESLA_L4_24TH'
----------------------------------------------------------------------
Ran 13 tests in 0.003s
FAILED (errors=1)
```
Contributor guide
Assessment
This issue has not been assessed yet.