Add CI workflow for tests that requires pytorch CUDA.

[vanbasten23]

This PR adds a new CI workflow that build pytorch with CUDA enabled from source, build pytorch/xla with CUDA enabled from source, then run tests. The intention is to run tests that requires pytorch with CUDA.

In detail, this PR add 2 more jobs to .github/workflows/build_and_test.yml

1. build-torch-with-cuda-xla-with-cuda: build pytorch with cuda and build pytorch/xla (It's ok to build torch_xla with cuda enabled now because it takes 2537s to build pytorch and takes 352s to build torch_xla).
2. test-cuda-with-pytorch-cuda-enabled: only run the tests that requires pytorch CUDA.

[vanbasten23]

Note to myself: it seems that BUILD XLA CUDA plugin requires the env var USE_CUDA=0. Otherwise, it'll fails with an error

      [5,777 / 9,258] Compiling tsl/platform/default/logging.cc; 0s local, remote-cache ... (17 actions, 2 running)
      ERROR: /github/home/.cache/bazel/_bazel_root/197a057057a49e5811107144e2d78508/external/xla/xla/service/gpu/BUILD:4486:13: Compiling xla/service/gpu/buffer_comparator.cu.cc failed: (Exit 1): crosstool_wrapper_driver_is_not_gcc failed: error executing command (from target @xla//xla/service/gpu:buffer_comparator_kernel) external/local_config_cuda/crosstool/clang/bin/crosstool_wrapper_driver_is_not_gcc -MD -MF bazel-out/k8-opt/bin/external/xla/xla/service/gpu/_objs/buffer_comparator_kernel/buffer_comparator.cu.pic.d ... (remaining 65 arguments skipped)
      nvcc fatal   : Unsupported gpu architecture 'compute_35'
      Target @xla//xla/pjrt/c:pjrt_c_api_gpu_plugin.so failed to build
      Use --verbose_failures to see the command lines of failed build steps.
      INFO: Elapsed time: 40.997s, Critical Path: 1.96s
      INFO: 5806 processes: 949 remote cache hit, 4857 internal.
      FAILED: Build did NOT complete successfully
      FAILED: Build did NOT complete successfully
      INFO: Streaming build results to: https://source.cloud.google.com/results/invocations/bba75106-7be5-4ad8-9a83-b3dd6222c3e1
      Traceback (most recent call last):
        File "/usr/local/lib/python3.10/site-packages/pip/_vendor/pyproject_hooks/_in_process/_in_process.py", line 353, in <module>
          main()
        File "/usr/local/lib/python3.10/site-packages/pip/_vendor/pyproject_hooks/_in_process/_in_process.py", line 335, in main
          json_out['return_val'] = hook(**hook_input['kwargs'])
        File "/usr/local/lib/python3.10/site-packages/pip/_vendor/pyproject_hooks/_in_process/_in_process.py", line 118, in get_requires_for_build_wheel
          return hook(config_settings)
        File "/tmp/pip-build-env-e97mw0k_/overlay/lib/python3.10/site-packages/setuptools/build_meta.py", line 325, in get_requires_for_build_wheel
          return self._get_build_requires(config_settings, requirements=['wheel'])
        File "/tmp/pip-build-env-e97mw0k_/overlay/lib/python3.10/site-packages/setuptools/build_meta.py", line 295, in _get_build_requires
          self.run_setup()
        File "/tmp/pip-build-env-e97mw0k_/overlay/lib/python3.10/site-packages/setuptools/build_meta.py", line 311, in run_setup
          exec(code, locals())
        File "<string>", line 11, in <module>
        File "/__w/xla/xla/pytorch/xla/plugins/cuda/../../build_util.py", line 67, in bazel_build
          subprocess.check_call(bazel_argv, stdout=sys.stdout, stderr=sys.stderr)
        File "/usr/local/lib/python3.10/subprocess.py", line 369, in check_call
          raise CalledProcessError(retcode, cmd)
      subprocess.CalledProcessError: Command '['bazel', 'build', '@xla//xla/pjrt/c:pjrt_c_api_gpu_plugin.so', '--symlink_prefix=/__w/xla/xla/pytorch/xla/plugins/cuda/bazel-', '--config=remote_cache', '--jobs=16', '--config=cuda', '--remote_default_exec_properties=cache-silo-key=dev']' returned non-zero exit status 1.
      error: subprocess-exited-with-error

[vanbasten23]

Seems without installing the cuda plugin, the tests would fail with error https://gist.github.com/vanbasten23/7dd6ddeaad93843e57653990c43cf476

[vanbasten23]

remove

[will-cromar]

This setup is repetitive. I was already on the fence about encapsulating it in a new action (since it already appears the build action, test actions, and the docs push. If we add another copy, it really should be encapsulated so we don't have to update a bunch of places at once.

[will-cromar]

This setup is repetitive. I was already on the fence about encapsulating it in a new action (since it already appears the build action, test actions, and the docs push). If we add another copy that's more-or-less identical, it really should be encapsulated so we don't have to update a bunch of places at once.

Stepping back, is there a way to merge this with _test.yml? You would need to add a parameter for some of the test groups to install the torch CUDA build.

[vanbasten23]

I can see your point. Let me give it a try.

[will-cromar]

I know this goes against my philosophy of "put everything in ansible", but what if you just build the PyTorch CUDA wheel directly here? I don't think we should build multiple copies of torch_xla and torchvision.

Building PyTorch with CUDA support is only part of our CI workflow, and it will never be part of our release workflow. It's okay in my mind to just directly USE_CUDA=1 python setup.py bdist_wheel here and upload only the torch GPU wheel as an artifact.

The test workflow can then use the same torch-xla, torch-xla-cuda-plugin, and torchvision.

[bhavya01]

Thanks for working on this. Added myself as a reviewer as I also need this for Triton tests.

[vanbasten23]

close it in favor of #7140