Disable cuda within one single pytest test

I have a code that should work on both gpu and cpu. Therefore, I write a test (pytest) to test the scenario of a device with cpu only or with gpu disabled.

I tried to modify the CUDA_VISIBLE_DEVICES environment variables but it does not work out.

def test_training_without_gpu(config):
environ_buffer = dict(environ)

    environ['CUDA_DEVICE_ORDER'] = 'PCI_BUS_ID'
    environ['CUDA_VISIBLE_DEVICES'] = ''  # hide all CUDA devices

    environ['Test'] = 'False'

    assert not cuda.is_available()
    training(config)

    # recover the env variables
    environ.clear()
    environ.update(environ_buffer)

Would anyone please help? :confused:

Hi,

This environment variable will only have effect when cuda is initialized. If you modify it after cuda is initialized, it will have no effect.
I don’t think you can disable cuda like that. You will need to use a manual flag or a “device” that you give to your test and that will be first cpu then cuda.

Hi,

Thank you. But is there any way to globally force to disable cuda? This will make sure the training can only run on cpu. The test will be more independent.
Otherwise if we use a flag, the code may still run gpu somewhere inside by mistake.

You can use CUDA_VISIBLE_DEVICES to hide all cuda devices but that needs to be done before any cuda initialization (so to be safe I usually do it in the command line like CUDA_VISIBLE_DEVICES=2 python my_prog.py).
The thing is that you won’t be able to change that value dynamically during the program. Only the value at the moment of the cuda initialization will be read and used. This is how cuda works unfortunately and out of our reach.