def test_gpu_file_lock_locking(tmpdir):
# the second time we try to acquire a lock for the same device we should not succeed
gpu_id = 0
candidates = [gpu_id]
with utils.GpuFileLock(candidates, str(tmpdir)) as lock_inner:
assert lock_inner == 0
with utils.GpuFileLock(candidates, str(tmpdir)) as lock_outer:
assert lock_outer is None
def test_gpu_file_lock_permission_exception(tmpdir):
with pytest.raises(PermissionError):
tmpdir = tmpdir.mkdir("sub")
# remove permissions
tmpdir.chmod(0)
with utils.GpuFileLock([0], str(tmpdir)) as lock:
assert False, "We expect to raise an exception when aquiring the lock and never reach this code."
def test_aquire_gpus_1_locked(tmpdir, requested_device_ids, num_gpus_available,
expected):
gpu_1 = 1
with utils.GpuFileLock([gpu_1], str(tmpdir)) as lock:
with utils.acquire_gpus(
requested_device_ids,
lock_dir=str(tmpdir),
num_gpus_available=num_gpus_available) as acquired_gpus:
assert set(acquired_gpus) == set(expected)
def test_gpu_file_lock_permission_exception(tmpdir):
tmpdir = tmpdir.mkdir("sub")
existing_lock = tmpdir.join("sockeye.gpu0.lock")
# remove permissions
existing_lock.write("")
existing_lock.chmod(0)
with utils.GpuFileLock([0, 1], str(tmpdir)) as acquired_lock:
# We expect to ignore the file for which we do not have permission and acquire the other device instead
assert acquired_lock == 1
def test_gpu_file_lock_cleanup(tmpdir):
gpu_id = 0
candidates = [gpu_id]
# Test that the lock files get created and clean up
with utils.GpuFileLock(candidates, str(tmpdir)) as lock:
assert lock == gpu_id
assert tmpdir.join("sockeye.gpu0.lock").check(), "Lock file did not exist."
assert not tmpdir.join("sockeye.gpu1.lock").check(), "Unrelated lock file did exist"
assert not tmpdir.join("sockeye.gpu0.lock").check(), "Lock file was not cleaned up."
def test_gpu_file_lock_exception_propagation(tmpdir):
gpu_ids = [0]
# Test that exceptions are properly propagated
raised_exception = RuntimeError("This exception should be propagated properly.")
caught_exception = None
try:
with utils.GpuFileLock(gpu_ids, str(tmpdir)) as lock:
raise raised_exception
except Exception as e:
caught_exception = e
assert caught_exception is raised_exception