Merge pull request #387 from chainer-ci/bp-386-v13-fix-cuda-header

[backport] Install `nvidia-cuda-runtime-cu12` in verifier

verifier/Dockerfile.debian

@@ -61,5 +61,6 @@ ENV LD_LIBRARY_PATH="/opt/rocm/lib:${LD_LIBRARY_PATH}"
 # Workaround for bug specific in ROCm 4.3 (https://github.com/cupy/cupy/issues/6605)
 ENV LLVM_PATH="/opt/rocm/llvm"
 
+COPY setup_cuda_runtime_headers.py /
 COPY agent.py /
 ENTRYPOINT ["/agent.py"]

verifier/Dockerfile.el8

@@ -26,4 +26,5 @@ RUN [ -z "${system_packages}" ] || ( \
 ENV HOME /tmp
 
 COPY agent.py /
+COPY setup_cuda_runtime_headers.py /
 ENTRYPOINT ["/agent.py"]

verifier/Dockerfile.rhel

@@ -29,4 +29,5 @@ RUN [ -z "${system_packages}" ] || ( \
 ENV HOME /tmp
 
 COPY agent.py /
+COPY setup_cuda_runtime_headers.py /
 ENTRYPOINT ["/agent.py"]

verifier/agent.py

@@ -58,6 +58,11 @@ def main(self):
         ]
         self._run(*cmdline)
 
+        self._log('Installing CUDA Runtime headers (if necessary)...')
+        verifier_dir = os.path.abspath(os.path.dirname(sys.argv[0]))
+        cmdline = pycommand + [f'{verifier_dir}/setup_cuda_runtime_headers.py']
+        self._run(*cmdline)
+
         # Importing CuPy should not be emit warnings,
         # Raise on warning to to catch bugs of preload warnings, e.g.:
         # https://github.com/cupy/cupy/pull/4933

verifier/setup_cuda_runtime_headers.py

@@ -0,0 +1,31 @@
+import shlex
+import subprocess
+import sys
+
+import cupy
+
+
+def main():
+    if cupy.cuda.runtime.is_hip:
+        return
+
+    # Only install if CUDA 12.2+.
+    (major, minor) = cupy.cuda.nvrtc.getVersion()
+    if major == 11:
+        return
+    elif major == 12:
+        if minor < 2:
+            return
+    else:
+        assert False, f'Unsupported CUDA version: {major}.{minor}'
+
+    cmdline = [
+        sys.executable, '-m', 'pip', 'install', '--user',
+        f'nvidia-cuda-runtime-cu{major}=={major}.{minor}.*',
+    ]
+    print(f'Running: {shlex.join(cmdline)}')
+    subprocess.run(cmdline, check=True)
+
+
+if __name__ == '__main__':
+    main()