From c81198789be183e7e1eb288eb98dd16f65b57e44 Mon Sep 17 00:00:00 2001
From: brandon-b-miller <53796099+brandon-b-miller@users.noreply.github.com>
Date: Wed, 17 Jan 2024 13:17:35 -0600
Subject: [PATCH] Defer PTX file load to runtime (#13690)

This PR fixes an issue where cuDF fails to import on machines with no NVIDIA GPU present.

cc @shwina

Authors:
  - https://github.com/brandon-b-miller
  - Vyas Ramasubramani (https://github.com/vyasr)

Approvers:
  - Bradley Dice (https://github.com/bdice)
  - Ashwin Srinath (https://github.com/shwina)
  - Vyas Ramasubramani (https://github.com/vyasr)

URL: https://github.com/rapidsai/cudf/pull/13690
---
 python/cudf/cudf/core/udf/utils.py         | 16 +++++++++++-----
 python/cudf/cudf/tests/test_no_device.py   | 16 ++++++++++++++++
 python/cudf/cudf/tests/test_string_udfs.py |  6 ++++--
 3 files changed, 31 insertions(+), 7 deletions(-)
 create mode 100644 python/cudf/cudf/tests/test_no_device.py

diff --git a/python/cudf/cudf/core/udf/utils.py b/python/cudf/cudf/core/udf/utils.py
index bd57db6b620..12baf1ea6d1 100644
--- a/python/cudf/cudf/core/udf/utils.py
+++ b/python/cudf/cudf/core/udf/utils.py
@@ -1,5 +1,6 @@
 # Copyright (c) 2020-2024, NVIDIA CORPORATION.
 
+import functools
 import os
 from typing import Any, Callable, Dict
 
@@ -60,10 +61,15 @@
 precompiled: cachetools.LRUCache = cachetools.LRUCache(maxsize=32)
 launch_arg_getters: Dict[Any, Any] = {}
 
-_PTX_FILE = _get_ptx_file(
-    os.path.join(os.path.dirname(strings_udf.__file__), "..", "core", "udf"),
-    "shim_",
-)
+
+@functools.cache
+def _ptx_file():
+    return _get_ptx_file(
+        os.path.join(
+            os.path.dirname(strings_udf.__file__), "..", "core", "udf"
+        ),
+        "shim_",
+    )
 
 
 @_cudf_nvtx_annotate
@@ -286,7 +292,7 @@ def _get_kernel(kernel_string, globals_, sig, func):
     exec(kernel_string, globals_)
     _kernel = globals_["_kernel"]
     kernel = cuda.jit(
-        sig, link=[_PTX_FILE], extensions=[str_view_arg_handler]
+        sig, link=[_ptx_file()], extensions=[str_view_arg_handler]
     )(_kernel)
 
     return kernel
diff --git a/python/cudf/cudf/tests/test_no_device.py b/python/cudf/cudf/tests/test_no_device.py
new file mode 100644
index 00000000000..722762b2d0c
--- /dev/null
+++ b/python/cudf/cudf/tests/test_no_device.py
@@ -0,0 +1,16 @@
+# Copyright (c) 2024, NVIDIA CORPORATION.
+import os
+import subprocess
+
+
+def test_cudf_import_no_device():
+    env = os.environ.copy()
+    env["CUDA_VISIBLE_DEVICES"] = "-1"
+    output = subprocess.run(
+        ["python", "-c", "import cudf"],
+        env=env,
+        capture_output=True,
+        text=True,
+        cwd="/",
+    )
+    assert output.returncode == 0
diff --git a/python/cudf/cudf/tests/test_string_udfs.py b/python/cudf/cudf/tests/test_string_udfs.py
index 88c73ccf964..5dbb86fe27d 100644
--- a/python/cudf/cudf/tests/test_string_udfs.py
+++ b/python/cudf/cudf/tests/test_string_udfs.py
@@ -1,4 +1,4 @@
-# Copyright (c) 2022-2023, NVIDIA CORPORATION.
+# Copyright (c) 2022-2024, NVIDIA CORPORATION.
 
 import numba
 import numpy as np
@@ -20,10 +20,12 @@
     string_view,
     udf_string,
 )
-from cudf.core.udf.utils import _PTX_FILE, _get_extensionty_size
+from cudf.core.udf.utils import _get_extensionty_size, _ptx_file
 from cudf.testing._utils import assert_eq, sv_to_udf_str
 from cudf.utils._numba import _CUDFNumbaConfig
 
+_PTX_FILE = _ptx_file()
+
 
 def get_kernels(func, dtype, size):
     """