[CI][Test] Fix uvm_test.py when a machine has only 1 GPU

[shintaro-iwasaki]

The following skipIf does not work correctly when a machine has only one GPU.

@unittest.skipIf(*gpu_unavailable or torch.cuda.device_count() < 2)
def test_uvm_to_device(self, sizes: List[int], uvm_op) -> None:
  [...]

This patch corrects the condition.

Details

First, *gpu_unavailable is defined as follows.

gpu_unavailable: Tuple[bool, str] = (
    not torch.cuda.is_available() or torch.cuda.device_count() == 0,
    "CUDA is not available or no GPUs detected",
)

So the skipIf is expanded as follows when there is only one CUDA device.

@unittest.skipIf(*gpu_unavailable or torch.cuda.device_count() < 2)
->
@unittest.skipIf((False, "CUDA is not ...") or True)

Because (False, "abc") or True is (False, "abc") in Python,

@unittest.skipIf((False, "CUDA is not ...") or True)
->
@unittest.skipIf(False, "CUDA is not ...")

It is False, so this unit test is not skipped. This UVM test seems failing occasionally because the machine does not have two GPUs, which annoys FBGEMM developers.

[netlify]

✅ Deploy Preview for pytorch-fbgemm-docs canceled.

Name	Link
🔨 Latest commit	138fc7b
🔍 Latest deploy log	https://app.netlify.com/sites/pytorch-fbgemm-docs/deploys/63c86937e0dca700094a05e3

[shintaro-iwasaki]

Now we see fbgemm_gpu/test/uvm_test.py::UvmTest::test_uvm_to_device SKIPPED on a single-GPU machine.

[facebook-github-bot]

@shintaro-iwasaki has imported this pull request. If you are a Meta employee, you can view this diff on Phabricator.

[facebook-github-bot]

@shintaro-iwasaki merged this pull request in 0237a8a.