diff --git a/ci/gpu/build.sh b/ci/gpu/build.sh index 4d86d19a7d2..7d67efa77b1 100755 --- a/ci/gpu/build.sh +++ b/ci/gpu/build.sh @@ -295,22 +295,15 @@ py.test -n 8 --cache-clear --basetemp="$WORKSPACE/custreamz-cuda-tmp" --junitxml gpuci_logger "Installing strings_udf" gpuci_mamba_retry install strings_udf -c "${CONDA_BLD_DIR}" -c "${CONDA_ARTIFACT_PATH}" -# only install strings_udf after cuDF is finished testing without its presence cd "$WORKSPACE/python/strings_udf/strings_udf" gpuci_logger "Python py.test for strings_udf" +py.test -n 8 --cache-clear --basetemp="$WORKSPACE/strings-udf-cuda-tmp" --junitxml="$WORKSPACE/junit-strings-udf.xml" -v --cov-config=.coveragerc --cov=strings_udf --cov-report=xml:"$WORKSPACE/python/strings_udf/strings-udf-coverage.xml" --cov-report term tests -STRINGS_UDF_PYTEST_RETCODE=0 -py.test -n 8 --cache-clear --basetemp="$WORKSPACE/strings-udf-cuda-tmp" --junitxml="$WORKSPACE/junit-strings-udf.xml" -v --cov-config=.coveragerc --cov=strings_udf --cov-report=xml:"$WORKSPACE/python/strings_udf/strings-udf-coverage.xml" --cov-report term tests || STRINGS_UDF_PYTEST_RETCODE=$? +# retest cuDF UDFs +cd "$WORKSPACE/python/cudf/cudf" +gpuci_logger "Python py.test retest cuDF UDFs" +py.test -n 8 --cache-clear --basetemp="$WORKSPACE/cudf-cuda-strings-udf-tmp" --ignore="$WORKSPACE/python/cudf/cudf/benchmarks" --junitxml="$WORKSPACE/junit-cudf-strings-udf.xml" -v --cov-config="$WORKSPACE/python/cudf/.coveragerc" --cov=cudf --cov-report=xml:"$WORKSPACE/python/cudf/cudf-strings-udf-coverage.xml" --cov-report term --dist=loadscope tests/test_udf_masked_ops.py -if [ ${STRINGS_UDF_PYTEST_RETCODE} -eq 5 ]; then - echo "No strings UDF tests were run, but this script will continue to execute." -elif [ ${STRINGS_UDF_PYTEST_RETCODE} -ne 0 ]; then - exit ${STRINGS_UDF_PYTEST_RETCODE} -else - cd "$WORKSPACE/python/cudf/cudf" - gpuci_logger "Python py.test retest cuDF UDFs" - py.test -n 8 --cache-clear --basetemp="$WORKSPACE/cudf-cuda-strings-udf-tmp" --ignore="$WORKSPACE/python/cudf/cudf/benchmarks" --junitxml="$WORKSPACE/junit-cudf-strings-udf.xml" -v --cov-config="$WORKSPACE/python/cudf/.coveragerc" --cov=cudf --cov-report=xml:"$WORKSPACE/python/cudf/cudf-strings-udf-coverage.xml" --cov-report term --dist=loadscope tests -fi # Run benchmarks with both cudf and pandas to ensure compatibility is maintained. # Benchmarks are run in DEBUG_ONLY mode, meaning that only small data sizes are used. diff --git a/conda/recipes/strings_udf/meta.yaml b/conda/recipes/strings_udf/meta.yaml index a736edef24d..027a8a82aae 100644 --- a/conda/recipes/strings_udf/meta.yaml +++ b/conda/recipes/strings_udf/meta.yaml @@ -50,7 +50,7 @@ requirements: - cudf ={{ version }} - {{ pin_compatible('cudatoolkit', max_pin='x', min_pin='x') }} - cachetools - - ptxcompiler # [linux64] # CUDA enhanced compatibility. See https://github.com/rapidsai/ptxcompiler + - ptxcompiler >=0.7.0 # [linux64] # CUDA enhanced compatibility. See https://github.com/rapidsai/ptxcompiler test: # [linux64] requires: # [linux64] - cudatoolkit {{ cuda_version }}.* # [linux64] diff --git a/python/cudf/cudf/core/udf/__init__.py b/python/cudf/cudf/core/udf/__init__.py index 443466b28bd..4730f1fa296 100644 --- a/python/cudf/cudf/core/udf/__init__.py +++ b/python/cudf/cudf/core/udf/__init__.py @@ -23,17 +23,20 @@ | {types.boolean} ) _STRING_UDFS_ENABLED = False + try: import strings_udf + from strings_udf import ptxpath + + if ptxpath: + utils.ptx_files.append(ptxpath) - if strings_udf.ENABLED: - from . import strings_typing # isort: skip - from . import strings_lowering # isort: skip - from strings_udf import ptxpath from strings_udf._lib.cudf_jit_udf import to_string_view_array from strings_udf._typing import str_view_arg_handler, string_view - # add an overload of MaskedType.__init__(string_view, bool) + from . import strings_typing # isort: skip + from . import strings_lowering # isort: skip + cuda_lower(api.Masked, strings_typing.string_view, types.boolean)( masked_lowering.masked_constructor ) @@ -47,13 +50,11 @@ utils.launch_arg_getters[dtype("O")] = to_string_view_array utils.masked_array_types[dtype("O")] = string_view utils.JIT_SUPPORTED_TYPES |= STRING_TYPES - utils.ptx_files.append(ptxpath) + utils.arg_handlers.append(str_view_arg_handler) row_function.itemsizes[dtype("O")] = string_view.size_bytes _STRING_UDFS_ENABLED = True - else: - del strings_udf except ImportError as e: # allow cuDF to work without strings_udf diff --git a/python/cudf/cudf/core/udf/strings_utils.py b/python/cudf/cudf/core/udf/strings_utils.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/python/strings_udf/strings_udf/__init__.py b/python/strings_udf/strings_udf/__init__.py index 2cedc0288d1..24f1a2d3bda 100644 --- a/python/strings_udf/strings_udf/__init__.py +++ b/python/strings_udf/strings_udf/__init__.py @@ -1,36 +1,20 @@ # Copyright (c) 2022, NVIDIA CORPORATION. import glob import os -import re -import subprocess -import sys +from cubinlinker.patch import _numba_version_ok, get_logger, new_patched_linker from numba import cuda -from ptxcompiler.patch import CMD +from numba.cuda.cudadrv.driver import Linker +from ptxcompiler.patch import NO_DRIVER, safe_get_versions from . import _version __version__ = _version.get_versions()["version"] -ENABLED = False +logger = get_logger() - -def compiler_from_ptx_file(path): - """Parse a PTX file header and extract the CUDA version used to compile it. - Here is an example PTX header that this function should parse: - // Generated by NVIDIA NVVM Compiler - // - // Compiler Build ID: CL-30672275 - // Cuda compilation tools, release 11.5, V11.5.119 - // Based on NVVM 7 - """ - file = open(path).read() - major, minor = ( - re.search(r"Cuda compilation tools, release ([0-9\.]+)", file) - .group(1) - .split(".") - ) - return int(major), int(minor) +# tracks the version of CUDA used to build the c++ and PTX components +STRINGS_UDF_PTX_VERSION = (11, 5) def _get_appropriate_file(sms, cc): @@ -41,59 +25,70 @@ def _get_appropriate_file(sms, cc): return None -# adapted from PTXCompiler -cp = subprocess.run([sys.executable, "-c", CMD], capture_output=True) -if cp.returncode == 0: - # must have a driver to proceed - versions = [int(s) for s in cp.stdout.strip().split()] - driver_version = tuple(versions[:2]) - runtime_version = tuple(versions[2:]) +def maybe_patch_numba_linker(driver_version): + # Numba thinks cubinlinker is only needed if the driver is older than the ctk + # but when strings_udf is present, it might also need to patch because the PTX + # file strings_udf relies on may be newer than the driver as well + if driver_version < STRINGS_UDF_PTX_VERSION: + logger.debug( + "Driver version %s.%s needs patching due to strings_udf" + % driver_version + ) + if _numba_version_ok: + logger.debug("Patching Numba Linker") + Linker.new = new_patched_linker + else: + logger.debug("Cannot patch Numba Linker - unsupported version") + + +def _get_ptx_file(): + if "RAPIDS_NO_INITIALIZE" in os.environ: + cc = int(os.environ.get("STRINGS_UDF_CC", "52")) + else: + dev = cuda.get_current_device() - # CUDA enhanced compatibility not yet enabled - if driver_version >= runtime_version: # Load the highest compute capability file available that is less than # the current device's. - dev = cuda.get_current_device() cc = int("".join(str(x) for x in dev.compute_capability)) - files = glob.glob( - os.path.join(os.path.dirname(__file__), "shim_*.ptx") + files = glob.glob(os.path.join(os.path.dirname(__file__), "shim_*.ptx")) + if len(files) == 0: + raise RuntimeError( + "This strings_udf installation is missing the necessary PTX " + f"files for compute capability {cc}. " + "Please file an issue reporting this error and how you " + "installed cudf and strings_udf." + "https://github.com/rapidsai/cudf/issues" ) - if len(files) == 0: - raise RuntimeError( - "This strings_udf installation is missing the necessary PTX " - "files. Please file an issue reporting this error and how you " - "installed cudf and strings_udf." - ) - - suffix_a_sm = None - regular_sms = [] - - for f in files: - file_name = os.path.basename(f) - sm_number = file_name.rstrip(".ptx").lstrip("shim_") - if sm_number.endswith("a"): - processed_sm_number = int(sm_number.rstrip("a")) - if processed_sm_number == cc: - suffix_a_sm = (processed_sm_number, f) - else: - regular_sms.append((int(sm_number), f)) - - regular_result = None - - if regular_sms: - regular_result = _get_appropriate_file(regular_sms, cc) - - if suffix_a_sm is None and regular_result is None: - raise RuntimeError( - "This strings_udf installation is missing the necessary PTX " - f"files that are <={cc}." - ) - elif suffix_a_sm is not None: - ptxpath = suffix_a_sm[1] - else: - ptxpath = regular_result[1] - if driver_version >= compiler_from_ptx_file(ptxpath): - ENABLED = True + regular_sms = [] + + for f in files: + file_name = os.path.basename(f) + sm_number = file_name.rstrip(".ptx").lstrip("shim_") + if sm_number.endswith("a"): + processed_sm_number = int(sm_number.rstrip("a")) + if processed_sm_number == cc: + return f else: - del ptxpath + regular_sms.append((int(sm_number), f)) + + regular_result = None + + if regular_sms: + regular_result = _get_appropriate_file(regular_sms, cc) + + if regular_result is None: + raise RuntimeError( + "This strings_udf installation is missing the necessary PTX " + f"files that are <={cc}." + ) + else: + return regular_result[1] + + +ptxpath = None +versions = safe_get_versions() +if versions != NO_DRIVER: + driver_version, runtime_version = versions + maybe_patch_numba_linker(driver_version) + ptxpath = _get_ptx_file() diff --git a/python/strings_udf/strings_udf/tests/test_string_udfs.py b/python/strings_udf/strings_udf/tests/test_string_udfs.py index f214915ae12..1a5dfa00dd7 100644 --- a/python/strings_udf/strings_udf/tests/test_string_udfs.py +++ b/python/strings_udf/strings_udf/tests/test_string_udfs.py @@ -15,9 +15,6 @@ from strings_udf._lib.cudf_jit_udf import to_string_view_array from strings_udf._typing import str_view_arg_handler, string_view -if not strings_udf.ENABLED: - pytest.skip("Strings UDF not enabled.", allow_module_level=True) - def get_kernel(func, dtype): """