def test_grabgpus3():
pop_if_set()
res = grab_gpus(100)
assert len(os.environ['CUDA_VISIBLE_DEVICES']) < len(','.join(
[str(x) for x in range(100)]))
nvmlInit()
assert res <= nvmlDeviceGetCount()
nvmlShutdown()
def test_grabgpus4():
res = grab_gpus(5, gpu_select=range(3, 8))
assert '0' not in os.environ['CUDA_VISIBLE_DEVICES']
assert '1' not in os.environ['CUDA_VISIBLE_DEVICES']
assert '2' not in os.environ['CUDA_VISIBLE_DEVICES']
def test_grabgpus3():
res = grab_gpus(100)
assert len(os.environ['CUDA_VISIBLE_DEVICES']) < len(','.join(
[str(x) for x in range(100)]))
assert res <= 8
def test_grabgpus2():
res = grab_gpus(1)
assert len(os.environ['CUDA_VISIBLE_DEVICES']) > 0
assert res == 1
def test_grabgpus():
res = grab_gpus(0)
assert os.environ['CUDA_VISIBLE_DEVICES'] == ''
assert res == 0
def test_called_twice():
pop_if_set()
grab_gpus(1)
grab_gpus(1, env_set_ok=True)
with pytest.raises(ValueError):
grab_gpus(1, env_set_ok=False)
