Skip to content

Commit

Permalink
Merge pull request #1 from brandon-b-miller/strings-udf-remove-cuda-r…
Browse files Browse the repository at this point in the history
…untime-checks

Remove runtime checks for CUDA versions from strings_udf
  • Loading branch information
gmarkall authored Sep 28, 2022
2 parents d9ad59b + 7374655 commit 1c344e8
Show file tree
Hide file tree
Showing 4 changed files with 33 additions and 119 deletions.
23 changes: 5 additions & 18 deletions ci/gpu/build.sh
Original file line number Diff line number Diff line change
Expand Up @@ -268,27 +268,14 @@ 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"

# We do not want to exit with a nonzero exit code in the case where no
# strings_udf tests are run because that will always happen when the local CUDA
# version is not 11.5. We need to suppress the exit code because this script is
# run with set -e and we're already setting a trap that we don't want to
# override here.

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=$?

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 tests/test_udf_masked_ops.py -n 8 --cache-clear
fi
# retest cudf with strings_udf present
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
gpuci_logger "Python py.test retest cuDF UDFs"
py.test tests/test_udf_masked_ops.py -n 8 --cache-clear

# 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.
Expand Down
64 changes: 28 additions & 36 deletions python/cudf/cudf/core/udf/__init__.py
Original file line number Diff line number Diff line change
@@ -1,10 +1,6 @@
# Copyright (c) 2022, NVIDIA CORPORATION.
import numpy as np
from numba import cuda, types
from numba.cuda.cudaimpl import (
lower as cuda_lower,
registry as cuda_lowering_registry,
)
from numba import types
from numba.cuda.cudaimpl import lower as cuda_lower

from cudf.core.dtypes import dtype
from cudf.core.udf import api, row_function, utils
Expand All @@ -28,36 +24,32 @@
_STRING_UDFS_ENABLED = False
try:
import strings_udf

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)
cuda_lower(api.Masked, strings_typing.string_view, types.boolean)(
masked_lowering.masked_constructor
)

# add an overload of pack_return(string_view)
cuda_lower(api.pack_return, strings_typing.string_view)(
masked_lowering.pack_return_scalar_impl
)

_supported_masked_types |= {strings_typing.string_view}
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

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

from . import strings_typing # isort: skip
from . import strings_lowering # isort: skip

# add an overload of MaskedType.__init__(string_view, bool)
cuda_lower(api.Masked, strings_typing.string_view, types.boolean)(
masked_lowering.masked_constructor
)

# add an overload of pack_return(string_view)
cuda_lower(api.pack_return, strings_typing.string_view)(
masked_lowering.pack_return_scalar_impl
)

_supported_masked_types |= {strings_typing.string_view}
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
except ImportError as e:
# allow cuDF to work without strings_udf
pass
Expand Down
62 changes: 0 additions & 62 deletions python/strings_udf/strings_udf/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -10,66 +10,4 @@

from . import _version

ENABLED = False


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)


# 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:])

# 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.
files = glob.glob(
os.path.join(os.path.dirname(__file__), "shim_*.ptx")
)
dev = cuda.get_current_device()
cc = "".join(str(x) for x in dev.compute_capability)
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 "
"files. Please file an issue reporting this error and how you "
"installed cudf and strings_udf."
)
sms = [
os.path.basename(f).rstrip(".ptx").lstrip("shim_") for f in files
]
selected_sm = max(sm for sm in sms if sm < cc)
ptxpath = os.path.join(
os.path.dirname(__file__), f"shim_{selected_sm}.ptx"
)

if driver_version >= compiler_from_ptx_file(ptxpath):
ENABLED = True
else:
del ptxpath

__version__ = _version.get_versions()["version"]
3 changes: 0 additions & 3 deletions python/strings_udf/strings_udf/tests/test_string_udfs.py
Original file line number Diff line number Diff line change
Expand Up @@ -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):
"""
Expand Down

0 comments on commit 1c344e8

Please sign in to comment.