Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

[PJRT|TPU] Update test_xla_devices_single_process_all_chips for expected device number #5421

Merged
merged 10 commits into from
Aug 9, 2023
Merged

[PJRT|TPU] Update test_xla_devices_single_process_all_chips for expected device number #5421

Changes from 9 commits
Commits
Show all changes
10 commits
Select commit Hold shift + click to select a range
8077bb1
Update test_runtime_tpu.py
ManfeiBai Aug 8, 2023
f1fa824
Update test_runtime_tpu.py
ManfeiBai Aug 8, 2023
aa66476
Update test_runtime_tpu.py
ManfeiBai Aug 8, 2023
d49ee9f
Update test_runtime_tpu.py
ManfeiBai Aug 8, 2023
36a7e89
Update test_runtime_tpu.py
ManfeiBai Aug 8, 2023
d20b283
Update test_runtime_tpu.py
ManfeiBai Aug 8, 2023
0d52e95
Update test_runtime_tpu.py
ManfeiBai Aug 8, 2023
58babb6
format
ManfeiBai Aug 8, 2023
bacc8b6
Update test_runtime_tpu.py to skip above 4
ManfeiBai Aug 9, 2023
fe115e4
Update test_runtime_tpu.py
ManfeiBai Aug 9, 2023
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
1 change: 1 addition & 0 deletions test/pjrt/test_runtime_tpu.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -92,6 +92,7 @@ def test_xla_devices_multiprocess(self):
devices_per_process = pjrt.run_multiprocess(xm.xla_device)
self.assertDictEqual(devices_per_process, expected)

@absltest.skipIf(tpu.num_available_devices() > 4, "Not implemented")
Copy link
Collaborator

Choose a reason for hiding this comment

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

Please make this tpu.num_available_chips. This test case is still valid when there are multiple devices per chip. Also, it's safer to do tpu.num_available_chips() != 4 here. I don't know the behavior of TPU_VISIBLE_CHIPS=0,1,2,3 when there are fewer than 4 chips actually available.

Copy link
Collaborator Author

Choose a reason for hiding this comment

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

Thanks, Will, updated

def test_xla_devices_single_process_all_chips(self):
will-cromar marked this conversation as resolved.
Show resolved Hide resolved
expected = _ordinal_to_device(
processes=1, cores_per_process=tpu.num_available_devices())
Expand Down