A CLI tool for checking if GPUs are available before running your script that uses GPUs.

pip install knock-on-gpus

You can use knock-on-gpus to run a script that uses GPUs.

knock-on-gpus -- python my_script.py

If some GPUs are not available, knock-on-gpus will return an error code and print a message to the console.

You can also use knock-on-gpus to run a script with specific GPUs.

CUDA_VISIBLE_DEVICES=0,1 knock-on-gpus -- python my_script.py

You can also use --devices or -d to specify the GPUs to use.

knock-on-gpus -d 0,1 -- python my_script.py

You can use --auto-select to automatically allocate the number of GPUs.

knock-on-gpus --auto-select 2 -- python my_script.py

If GPU:0, GPU:1, and GPU:3 are unavailable, knock-on-gpus will use GPU:2 and GPU:4.

You can set an alias for python to use knock-on-gpus by default.

alias unsafe-python="`which python`"
alias python="knock-on-gpus -- python"

Then you can run your script without knock-on-gpus.

(Alias: -d, --device)

Specifies the GPUs to use. The value is a comma-separated list of GPU IDs.

Specifies the memory border (MiB) to treat as vacant. If the memory usage exceeds this value, the GPU will be treated as occupied.

If true, use GPU strictly. If CUDA is not available, it will fail.

Specifies the number of min GPUs to use.

Specifies the number of max GPUs to use.

Specifies the environment variable key to set visible devices.

If true, print verbose logs.

(Alias: -a, --auto)

If a number is given, it will automatically allocate the number of GPUs.

If true, allow running without executing extra commands.

$ knock-on-gpus -d "0,1,2,3" -- sh -c 'echo "devices=$CUDA_VISIBLE_DEVICES"'
# => devices=0,1,2,3

If GPUs are available, this will succeed. "devices=0,1,2,3" will be printed.

$ knock-on-gpus && sh -c 'echo "devices=$CUDA_VISIBLE_DEVICES"'
# => ERROR: Omitting the command is not allowed.

Even if GPUs are available, this will fail because no command passed to knock-on-gpus.

Note that && has no effect to pass the command to knock-on-gpus.

$ knock-on-gpus --allow-noop && sh -c 'echo "devices=$CUDA_VISIBLE_DEVICES"'
# => devices=

If GPUs are available, this will succeed. BUT "devices=" is printed instead of "devices=0,1,2,3" because sh -c ... is not passed to knock-on-gpus.