Unable to use nvdiffrast with Sagemaker Distribution 2.2.1

### Product Category

JupyterLab

### Question

Im trying to use [nvdiffrast](https://github.com/NVlabs/nvdiffrast) on jupyterlab with Sagemaker Distribution 2.2.1 image and ml.g5.2xlarge instance.
I can successfully install it, but when running a test code
```
import nvdiffrast
import nvdiffrast.torch as dr

a = dr.RasterizeCudaContext(device='cuda:0')
```
I get an error
```
/opt/conda/lib/python3.11/site-packages/torch/utils/cpp_extension.py:1967: UserWarning: TORCH_CUDA_ARCH_LIST is not set, all archs for visible cards are included for compilation.
If this is not desired, please set os.environ['TORCH_CUDA_ARCH_LIST'].
warnings.warn(
Traceback (most recent call last):
File "/opt/conda/lib/python3.11/site-packages/torch/utils/cpp_extension.py", line 2107, in _run_ninja_build
subprocess.run(
File "/opt/conda/lib/python3.11/subprocess.py", line 571, in run
raise CalledProcessError(retcode, process.args,
subprocess.CalledProcessError: Command '['ninja', '-v']' returned non-zero exit status 1.

The above exception was the direct cause of the following exception:

Traceback (most recent call last):
File "/home/sagemaker-user/Wonder3D/t.py", line 7, in
a = dr.RasterizeCudaContext(device='cuda')
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/opt/conda/lib/python3.11/site-packages/nvdiffrast/torch/ops.py", line 184, in __init__
self.cpp_wrapper = _get_plugin().RasterizeCRStateWrapper(cuda_device_idx)
^^^^^^^^^^^^^
File "/opt/conda/lib/python3.11/site-packages/nvdiffrast/torch/ops.py", line 125, in _get_plugin
torch.utils.cpp_extension.load(name=plugin_name, sources=source_paths, extra_cflags=common_opts+cc_opts, extra_cuda_cflags=common_opts+['-lineinfo'], extra_ldflags=ldflags, with_cuda=True, verbose=False)
File "/opt/conda/lib/python3.11/site-packages/torch/utils/cpp_extension.py", line 1309, in load
return _jit_compile(
^^^^^^^^^^^^^
File "/opt/conda/lib/python3.11/site-packages/torch/utils/cpp_extension.py", line 1719, in _jit_compile
_write_ninja_file_and_build_library(
File "/opt/conda/lib/python3.11/site-packages/torch/utils/cpp_extension.py", line 1832, in _write_ninja_file_and_build_library
_run_ninja_build(
File "/opt/conda/lib/python3.11/site-packages/torch/utils/cpp_extension.py", line 2123, in _run_ninja_build
raise RuntimeError(message) from e
RuntimeError: Error building extension 'nvdiffrast_plugin': [1/11] /opt/conda/bin/x86_64-conda-linux-gnu-c++ -MMD -MF CudaRaster.o.d -DTORCH_EXTENSION_NAME=nvdiffrast_plugin -DTORCH_API_INCLUDE_EXTENSION_H -DPYBIND11_COMPILER_TYPE=\"_gcc\" -DPYBIND11_STDLIB=\"_libstdcpp\" -DPYBIND11_BUILD_ABI=\"_cxxabi1017\" -isystem /opt/conda/lib/python3.11/site-packages/torch/include -isystem /opt/conda/lib/python3.11/site-packages/torch/include/torch/csrc/api/include -isystem /opt/conda/lib/python3.11/site-packages/torch/include/TH -isystem /opt/conda/lib/python3.11/site-packages/torch/include/THC -isystem /opt/conda/include -isystem /opt/conda/include/python3.11 -D_GLIBCXX_USE_CXX11_ABI=1 -fPIC -std=c++17 -DNVDR_TORCH -c /opt/conda/lib/python3.11/site-packages/nvdiffrast/common/cudaraster/impl/CudaRaster.cpp -o CudaRaster.o
FAILED: CudaRaster.o
/opt/conda/bin/x86_64-conda-linux-gnu-c++ -MMD -MF CudaRaster.o.d -DTORCH_EXTENSION_NAME=nvdiffrast_plugin -DTORCH_API_INCLUDE_EXTENSION_H -DPYBIND11_COMPILER_TYPE=\"_gcc\" -DPYBIND11_STDLIB=\"_libstdcpp\" -DPYBIND11_BUILD_ABI=\"_cxxabi1017\" -isystem /opt/conda/lib/python3.11/site-packages/torch/include -isystem /opt/conda/lib/python3.11/site-packages/torch/include/torch/csrc/api/include -isystem /opt/conda/lib/python3.11/site-packages/torch/include/TH -isystem /opt/conda/lib/python3.11/site-packages/torch/include/THC -isystem /opt/conda/include -isystem /opt/conda/include/python3.11 -D_GLIBCXX_USE_CXX11_ABI=1 -fPIC -std=c++17 -DNVDR_TORCH -c /opt/conda/lib/python3.11/site-packages/nvdiffrast/common/cudaraster/impl/CudaRaster.cpp -o CudaRaster.o
In file included from /opt/conda/lib/python3.11/site-packages/nvdiffrast/common/cudaraster/impl/CudaRaster.cpp:9:
/opt/conda/lib/python3.11/site-packages/nvdiffrast/common/cudaraster/impl/Defs.hpp:10:10: fatal error: cuda_runtime.h: No such file or directory
10 | #include
| ^~~~~~~~~~~~~~~~
compilation terminated.
[2/11] /opt/conda/bin/x86_64-conda-linux-gnu-c++ -MMD -MF common.o.d -DTORCH_EXTENSION_NAME=nvdiffrast_plugin -DTORCH_API_INCLUDE_EXTENSION_H -DPYBIND11_COMPILER_TYPE=\"_gcc\" -DPYBIND11_STDLIB=\"_libstdcpp\" -DPYBIND11_BUILD_ABI=\"_cxxabi1017\" -isystem /opt/conda/lib/python3.11/site-packages/torch/include -isystem /opt/conda/lib/python3.11/site-packages/torch/include/torch/csrc/api/include -isystem /opt/conda/lib/python3.11/site-packages/torch/include/TH -isystem /opt/conda/lib/python3.11/site-packages/torch/include/THC -isystem /opt/conda/include -isystem /opt/conda/include/python3.11 -D_GLIBCXX_USE_CXX11_ABI=1 -fPIC -std=c++17 -DNVDR_TORCH -c /opt/conda/lib/python3.11/site-packages/nvdiffrast/common/common.cpp -o common.o
FAILED: common.o
/opt/conda/bin/x86_64-conda-linux-gnu-c++ -MMD -MF common.o.d -DTORCH_EXTENSION_NAME=nvdiffrast_plugin -DTORCH_API_INCLUDE_EXTENSION_H -DPYBIND11_COMPILER_TYPE=\"_gcc\" -DPYBIND11_STDLIB=\"_libstdcpp\" -DPYBIND11_BUILD_ABI=\"_cxxabi1017\" -isystem /opt/conda/lib/python3.11/site-packages/torch/include -isystem /opt/conda/lib/python3.11/site-packages/torch/include/torch/csrc/api/include -isystem /opt/conda/lib/python3.11/site-packages/torch/include/TH -isystem /opt/conda/lib/python3.11/site-packages/torch/include/THC -isystem /opt/conda/include -isystem /opt/conda/include/python3.11 -D_GLIBCXX_USE_CXX11_ABI=1 -fPIC -std=c++17 -DNVDR_TORCH -c /opt/conda/lib/python3.11/site-packages/nvdiffrast/common/common.cpp -o common.o
/opt/conda/lib/python3.11/site-packages/nvdiffrast/common/common.cpp:9:10: fatal error: cuda_runtime.h: No such file or directory
9 | #include
| ^~~~~~~~~~~~~~~~
compilation terminated.
[3/11] /opt/conda/bin/x86_64-conda-linux-gnu-c++ -MMD -MF torch_bindings.o.d -DTORCH_EXTENSION_NAME=nvdiffrast_plugin -DTORCH_API_INCLUDE_EXTENSION_H -DPYBIND11_COMPILER_TYPE=\"_gcc\" -DPYBIND11_STDLIB=\"_libstdcpp\" -DPYBIND11_BUILD_ABI=\"_cxxabi1017\" -isystem /opt/conda/lib/python3.11/site-packages/torch/include -isystem /opt/conda/lib/python3.11/site-packages/torch/include/torch/csrc/api/include -isystem /opt/conda/lib/python3.11/site-packages/torch/include/TH -isystem /opt/conda/lib/python3.11/site-packages/torch/include/THC -isystem /opt/conda/include -isystem /opt/conda/include/python3.11 -D_GLIBCXX_USE_CXX11_ABI=1 -fPIC -std=c++17 -DNVDR_TORCH -c /opt/conda/lib/python3.11/site-packages/nvdiffrast/torch/torch_bindings.cpp -o torch_bindings.o
FAILED: torch_bindings.o
/opt/conda/bin/x86_64-conda-linux-gnu-c++ -MMD -MF torch_bindings.o.d -DTORCH_EXTENSION_NAME=nvdiffrast_plugin -DTORCH_API_INCLUDE_EXTENSION_H -DPYBIND11_COMPILER_TYPE=\"_gcc\" -DPYBIND11_STDLIB=\"_libstdcpp\" -DPYBIND11_BUILD_ABI=\"_cxxabi1017\" -isystem /opt/conda/lib/python3.11/site-packages/torch/include -isystem /opt/conda/lib/python3.11/site-packages/torch/include/torch/csrc/api/include -isystem /opt/conda/lib/python3.11/site-packages/torch/include/TH -isystem /opt/conda/lib/python3.11/site-packages/torch/include/THC -isystem /opt/conda/include -isystem /opt/conda/include/python3.11 -D_GLIBCXX_USE_CXX11_ABI=1 -fPIC -std=c++17 -DNVDR_TORCH -c /opt/conda/lib/python3.11/site-packages/nvdiffrast/torch/torch_bindings.cpp -o torch_bindings.o
In file included from /opt/conda/include/ATen/cuda/CUDAContext.h:3,
from /opt/conda/lib/python3.11/site-packages/nvdiffrast/torch/../common/framework.h:37,
from /opt/conda/lib/python3.11/site-packages/nvdiffrast/torch/torch_common.inl:10,
from /opt/conda/lib/python3.11/site-packages/nvdiffrast/torch/torch_bindings.cpp:9:
/opt/conda/include/ATen/cuda/CUDAContextLight.h:6:10: fatal error: cuda_runtime_api.h: No such file or directory
6 | #include
| ^~~~~~~~~~~~~~~~~~~~
compilation terminated.
[4/11] /opt/conda/bin/x86_64-conda-linux-gnu-c++ -MMD -MF torch_texture.o.d -DTORCH_EXTENSION_NAME=nvdiffrast_plugin -DTORCH_API_INCLUDE_EXTENSION_H -DPYBIND11_COMPILER_TYPE=\"_gcc\" -DPYBIND11_STDLIB=\"_libstdcpp\" -DPYBIND11_BUILD_ABI=\"_cxxabi1017\" -isystem /opt/conda/lib/python3.11/site-packages/torch/include -isystem /opt/conda/lib/python3.11/site-packages/torch/include/torch/csrc/api/include -isystem /opt/conda/lib/python3.11/site-packages/torch/include/TH -isystem /opt/conda/lib/python3.11/site-packages/torch/include/THC -isystem /opt/conda/include -isystem /opt/conda/include/python3.11 -D_GLIBCXX_USE_CXX11_ABI=1 -fPIC -std=c++17 -DNVDR_TORCH -c /opt/conda/lib/python3.11/site-packages/nvdiffrast/torch/torch_texture.cpp -o torch_texture.o
FAILED: torch_texture.o
/opt/conda/bin/x86_64-conda-linux-gnu-c++ -MMD -MF torch_texture.o.d -DTORCH_EXTENSION_NAME=nvdiffrast_plugin -DTORCH_API_INCLUDE_EXTENSION_H -DPYBIND11_COMPILER_TYPE=\"_gcc\" -DPYBIND11_STDLIB=\"_libstdcpp\" -DPYBIND11_BUILD_ABI=\"_cxxabi1017\" -isystem /opt/conda/lib/python3.11/site-packages/torch/include -isystem /opt/conda/lib/python3.11/site-packages/torch/include/torch/csrc/api/include -isystem /opt/conda/lib/python3.11/site-packages/torch/include/TH -isystem /opt/conda/lib/python3.11/site-packages/torch/include/THC -isystem /opt/conda/include -isystem /opt/conda/include/python3.11 -D_GLIBCXX_USE_CXX11_ABI=1 -fPIC -std=c++17 -DNVDR_TORCH -c /opt/conda/lib/python3.11/site-packages/nvdiffrast/torch/torch_texture.cpp -o torch_texture.o
In file included from /opt/conda/include/ATen/cuda/CUDAContext.h:3,
from /opt/conda/lib/python3.11/site-packages/nvdiffrast/torch/../common/framework.h:37,
from /opt/conda/lib/python3.11/site-packages/nvdiffrast/torch/torch_common.inl:10,
from /opt/conda/lib/python3.11/site-packages/nvdiffrast/torch/torch_texture.cpp:9:
/opt/conda/include/ATen/cuda/CUDAContextLight.h:6:10: fatal error: cuda_runtime_api.h: No such file or directory
6 | #include
| ^~~~~~~~~~~~~~~~~~~~
compilation terminated.
[5/11] /opt/conda/bin/x86_64-conda-linux-gnu-c++ -MMD -MF torch_interpolate.o.d -DTORCH_EXTENSION_NAME=nvdiffrast_plugin -DTORCH_API_INCLUDE_EXTENSION_H -DPYBIND11_COMPILER_TYPE=\"_gcc\" -DPYBIND11_STDLIB=\"_libstdcpp\" -DPYBIND11_BUILD_ABI=\"_cxxabi1017\" -isystem /opt/conda/lib/python3.11/site-packages/torch/include -isystem /opt/conda/lib/python3.11/site-packages/torch/include/torch/csrc/api/include -isystem /opt/conda/lib/python3.11/site-packages/torch/include/TH -isystem /opt/conda/lib/python3.11/site-packages/torch/include/THC -isystem /opt/conda/include -isystem /opt/conda/include/python3.11 -D_GLIBCXX_USE_CXX11_ABI=1 -fPIC -std=c++17 -DNVDR_TORCH -c /opt/conda/lib/python3.11/site-packages/nvdiffrast/torch/torch_interpolate.cpp -o torch_interpolate.o
```

This might indicate that there is no cuda, but the output of nvcc --version is:
```
nvcc: NVIDIA (R) Cuda compiler driver
Copyright (c) 2005-2024 NVIDIA Corporation
Built on Thu_Jun__6_02:18:23_PDT_2024
Cuda compilation tools, release 12.5, V12.5.82
Build cuda_12.5.r12.5/compiler.34385749_0

```

### Other Details

_No response_

Start with the reproduction in the issue and inspect nvdiffrast/torch/ops.py at _get_plugin(), then compare the available CUDA compiler and headers in SageMaker Distribution 2.2.1. Done means the RasterizeCudaContext example builds its extension successfully on the stated instance, or the supported environment limitation is documented.

Written by the indexing model from the issue text.