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]])