def testFailsWithBadInput(self):
with self.assertRaises(ValueError):
distributeGpuIdsToTasks([1, 2, 3], 5)
with self.assertRaises(ValueError):
distributeGpuIdsToTasks([1, 2, 3], 2)
with self.assertRaises(TypeError):
distributeGpuIdsToTasks(1, 5)
def testHandlesNonZeroStartingIDs(self):
self.assertEqual(distributeGpuIdsToTasks([5, 18], 1), [[5, 18]])
def testMultiTaskMultiGpu(self):
# more GPUs than tasks
self.assertEqual(distributeGpuIdsToTasks([0, 1, 2, 3], 2), [[0, 1], [2, 3]])
# more tasks than GPUs
self.assertEqual(distributeGpuIdsToTasks([0, 1], 6), [[0], [0], [0], [1], [1], [1]])
def testMultiTaskSingleGpu(self):
self.assertEqual(distributeGpuIdsToTasks([0], 3), [[0], [0], [0]])
def testSingleTaskMultipleGpu(self):
self.assertEqual(distributeGpuIdsToTasks([0, 1, 2], 1), [[0, 1, 2]])
def testSingleTaskSingleGpu(self):
self.assertEqual(distributeGpuIdsToTasks([0], 1), [[0]])