Skip to content

GitLab

"git@developer.sourcefind.cn:zhaoyu6/sglang.git" did not exist on "bed05878f6c8f3a7e730ddef450b728cbe5ad308"
Commit b2522f9b authored by josh11b's avatar josh11b Committed by Taylor Robie
Browse files

worker_devices moving to ds.extended (#5775)

parent 9a4848a2
Show whitespace changes
Inline Side-by-side
Showing with 6 additions and 6 deletions
official/utils/misc/distribution_utils_test.py
View file @ b2522f9b
......@@ -28,20 +28,20 @@ class GetDistributionStrategyTest(tf.test.TestCase):
def test_one_device_strategy_cpu(self):
ds = distribution_utils.get_distribution_strategy(0)
self.assertEquals(ds.num_replicas_in_sync, 1)
self.assertEquals(len(ds.worker_devices), 1)
self.assertIn('CPU', ds.worker_devices[0])
self.assertEquals(len(ds.extended.worker_devices), 1)
self.assertIn('CPU', ds.extended.worker_devices[0])
def test_one_device_strategy_gpu(self):
ds = distribution_utils.get_distribution_strategy(1)
self.assertEquals(ds.num_replicas_in_sync, 1)
self.assertEquals(len(ds.worker_devices), 1)
self.assertIn('GPU', ds.worker_devices[0])
self.assertEquals(len(ds.extended.worker_devices), 1)
self.assertIn('GPU', ds.extended.worker_devices[0])
def test_mirrored_strategy(self):
ds = distribution_utils.get_distribution_strategy(5)
self.assertEquals(ds.num_replicas_in_sync, 5)
self.assertEquals(len(ds.worker_devices), 5)
for device in ds.worker_devices:
self.assertEquals(len(ds.extended.worker_devices), 5)
for device in ds.extended.worker_devices:
self.assertIn('GPU', device)
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment