diff --git a/.github/workflows/native-lib-extraction-test.yml b/.github/workflows/native-lib-extraction-test.yml new file mode 100644 index 00000000..ae14fa7d --- /dev/null +++ b/.github/workflows/native-lib-extraction-test.yml @@ -0,0 +1,44 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +name: Native library extraction tests + +on: + push: + branches: + - main + paths: + - "examples/UDF-Examples/RAPIDS-accelerated-UDFs/extract-native-libs.py" + - "examples/UDF-Examples/RAPIDS-accelerated-UDFs/test_extract_native_libs.py" + - ".github/workflows/native-lib-extraction-test.yml" + pull_request: + paths: + - "examples/UDF-Examples/RAPIDS-accelerated-UDFs/extract-native-libs.py" + - "examples/UDF-Examples/RAPIDS-accelerated-UDFs/test_extract_native_libs.py" + - ".github/workflows/native-lib-extraction-test.yml" + +permissions: + contents: read + +jobs: + test: + runs-on: ubuntu-latest + steps: + - uses: actions/checkout@v4 + - name: Run native library extraction tests + run: | + python3 -m unittest discover \ + -s examples/UDF-Examples/RAPIDS-accelerated-UDFs \ + -p "test_extract_native_libs.py" \ + -v diff --git a/examples/UDF-Examples/RAPIDS-accelerated-UDFs/README.md b/examples/UDF-Examples/RAPIDS-accelerated-UDFs/README.md index 6fc5995f..6712ac3a 100644 --- a/examples/UDF-Examples/RAPIDS-accelerated-UDFs/README.md +++ b/examples/UDF-Examples/RAPIDS-accelerated-UDFs/README.md @@ -145,7 +145,8 @@ mvn clean package -Pudf-native-examples ``` The build will automatically: -- Extract `libcudf.so` from the rapids-4-spark jar +- Extract `libcudf.so` from older rapids-4-spark jars or reconstruct it from + the chunk-manifest representation used by newer jars - Read the embedded `spark-rapids-jni` and `cudf-java` version metadata from the jar - Download the matching `spark-rapids-jni` `cudf-pins` files - Clone the cuDF repository at the revision recorded in the jar diff --git a/examples/UDF-Examples/RAPIDS-accelerated-UDFs/extract-cudf-libs.sh b/examples/UDF-Examples/RAPIDS-accelerated-UDFs/extract-cudf-libs.sh index 942e2873..d4533a73 100755 --- a/examples/UDF-Examples/RAPIDS-accelerated-UDFs/extract-cudf-libs.sh +++ b/examples/UDF-Examples/RAPIDS-accelerated-UDFs/extract-cudf-libs.sh @@ -149,103 +149,25 @@ else exit 1 fi -# Extract libcudf.so and dependencies +# Extract or reconstruct libcudf.so and dependencies echo "Extracting native libraries from jar..." echo " Jar: $JAR_PATH" -echo " Looking for: */libcudf.so*, */libnvcomp.so*" +echo " Supports conventional and chunked native library entries" -# Use unzip without -q to capture output, but redirect to log for debugging -UNZIP_OUTPUT=$(unzip -o "$JAR_PATH" "*/libcudf.so*" "*/libnvcomp.so*" -d "$TARGET_DIR/temp" 2>&1) -UNZIP_EXIT_CODE=$? - -# Check unzip exit code -if [ $UNZIP_EXIT_CODE -ne 0 ]; then - echo "ERROR: Failed to extract libraries from jar" >&2 - echo "unzip exit code: $UNZIP_EXIT_CODE" >&2 - - # Provide helpful diagnostics - case $UNZIP_EXIT_CODE in - 11) - echo "Reason: No matching files found in jar" >&2 - echo "" >&2 - echo "The jar may not contain native libraries for your platform." >&2 - echo "Expected patterns: */libcudf.so*, */libnvcomp.so*" >&2 - echo "" >&2 - echo "Listing jar contents:" >&2 - unzip -l "$JAR_PATH" | grep -E '\.(so|dylib|dll)' || echo " No native libraries found" >&2 - ;; - *) - echo "Reason: unzip command failed" >&2 - echo "Output: $UNZIP_OUTPUT" >&2 - ;; - esac - - echo "" >&2 - echo "Falling back to source build..." >&2 - exit 1 -fi - -# Verify that we actually extracted some files -EXTRACTED_COUNT=$(find "$TARGET_DIR/temp" -name "*.so*" 2>/dev/null | wc -l) -echo "Extracted $EXTRACTED_COUNT library file(s)" - -if [ "$EXTRACTED_COUNT" -eq 0 ]; then - echo "ERROR: No library files were extracted from jar" >&2 - echo "This usually means the jar doesn't contain native libraries." >&2 - echo "" >&2 - echo "Listing jar contents:" >&2 - unzip -l "$JAR_PATH" | head -20 >&2 +if ! command -v python3 >/dev/null 2>&1; then + echo "ERROR: python3 is required to extract native libraries" >&2 exit 1 fi -# Move libraries to native-deps directory, detecting conflicts -echo "Moving extracted libraries..." -CONFLICT_COUNT=0 - -# Use process substitution to avoid subshell issues -while IFS= read -r source_file; do - filename=$(basename "$source_file") - dest_file="$NATIVE_DEPS_DIR/$filename" - - if [ -f "$dest_file" ]; then - # File already exists - check if it's the same - if ! cmp -s "$source_file" "$dest_file"; then - echo "WARNING: Conflicting library detected: $filename" >&2 - echo " Existing: $dest_file" >&2 - echo " New: $source_file" >&2 - echo " Keeping existing file, skipping new one" >&2 - CONFLICT_COUNT=$((CONFLICT_COUNT + 1)) - fi - # Remove the duplicate source file - rm -f "$source_file" - else - # No conflict, move the file - mv "$source_file" "$dest_file" - fi -done < <(find "$TARGET_DIR/temp" -name "*.so*") - -if [ "$CONFLICT_COUNT" -gt 0 ]; then - echo "WARNING: $CONFLICT_COUNT library file(s) had conflicts. Review the warnings above." >&2 -fi - -rm -rf "$TARGET_DIR/temp" +python3 "$SCRIPT_DIR/extract-native-libs.py" "$JAR_PATH" "$NATIVE_DEPS_DIR" -# Verify that libcudf.so was successfully moved to final location +# Verify that libcudf.so was successfully written to the final location if [ ! -f "$NATIVE_DEPS_DIR/libcudf.so" ]; then echo "ERROR: libcudf.so not found in $NATIVE_DEPS_DIR" >&2 - echo "" >&2 - echo "This could mean:" >&2 - echo " 1. The jar didn't contain libcudf.so" >&2 - echo " 2. Extraction succeeded but moving files failed" >&2 - echo " 3. Wrong architecture (jar might be for a different platform)" >&2 - echo "" >&2 - echo "Contents of $NATIVE_DEPS_DIR:" >&2 - ls -lh "$NATIVE_DEPS_DIR" >&2 || echo " Directory is empty or doesn't exist" >&2 exit 1 fi echo "✓ Successfully extracted libraries to: $NATIVE_DEPS_DIR" -ls -lh "$NATIVE_DEPS_DIR" PINS_DIR="$TARGET_DIR/cudf-pins" PINS_PROPERTIES="$TARGET_DIR/jar-native-deps.properties" diff --git a/examples/UDF-Examples/RAPIDS-accelerated-UDFs/extract-native-libs.py b/examples/UDF-Examples/RAPIDS-accelerated-UDFs/extract-native-libs.py new file mode 100644 index 00000000..cded5680 --- /dev/null +++ b/examples/UDF-Examples/RAPIDS-accelerated-UDFs/extract-native-libs.py @@ -0,0 +1,248 @@ +#!/usr/bin/env python3 +# +# Copyright (c) 2026, NVIDIA CORPORATION. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# + +"""Extract native cuDF libraries from a rapids-4-spark jar. + +Large native libraries may be stored as a versioned manifest and numbered +chunks. This script supports both that representation and conventional jar +entries such as amd64/Linux/libcudf.so. +""" + +import argparse +import binascii +import os +import shutil +import tempfile +import zipfile + + +COPY_BUFFER_SIZE = 1024 * 1024 +MANIFEST_SUFFIX = ".chunks.properties" +CHUNK_DIRECTORY_SUFFIX = ".chunks" +CUDF_LIBRARY_NAME = "libcudf.so" + + +def parse_properties(data, manifest_name): + properties = {} + for raw_line in data.decode("utf-8").splitlines(): + line = raw_line.strip() + if not line or line.startswith(("#", "!")): + continue + key, separator, value = line.partition("=") + if not separator: + raise RuntimeError(f"Invalid property in {manifest_name}: {raw_line!r}") + properties[key.strip()] = value.strip() + return properties + + +def require_property(properties, key, manifest_name): + value = properties.get(key) + if not value: + raise RuntimeError(f"Missing {key} in {manifest_name}") + return value + + +def parse_positive_int(properties, key, manifest_name): + value = require_property(properties, key, manifest_name) + try: + result = int(value) + except ValueError as error: + raise RuntimeError(f"Invalid {key} in {manifest_name}: {value}") from error + if result <= 0: + raise RuntimeError(f"{key} must be positive in {manifest_name}: {value}") + return result + + +def temporary_output(output_dir, library_name): + descriptor, path = tempfile.mkstemp(prefix=f".{library_name}.", dir=output_dir) + return os.fdopen(descriptor, "wb"), path + + +def install_output(temporary_path, output_path): + os.chmod(temporary_path, 0o755) + os.replace(temporary_path, output_path) + + +def extract_conventional_library(archive, entry_name, output_dir): + library_name = os.path.basename(entry_name) + output_path = os.path.join(output_dir, library_name) + output, temporary_path = temporary_output(output_dir, library_name) + try: + with output, archive.open(entry_name) as source: + shutil.copyfileobj(source, output, COPY_BUFFER_SIZE) + install_output(temporary_path, output_path) + except Exception: + try: + os.remove(temporary_path) + except FileNotFoundError: + pass + raise + print(f"Extracted {entry_name} -> {output_path}") + + +def reconstruct_chunked_library(archive, manifest_name, output_dir): + library_entry = manifest_name[: -len(MANIFEST_SUFFIX)] + library_name = os.path.basename(library_entry) + properties = parse_properties(archive.read(manifest_name), manifest_name) + + format_version = require_property(properties, "format.version", manifest_name) + if format_version != "1": + raise RuntimeError( + f"Unsupported native chunk format.version in {manifest_name}: {format_version}" + ) + + library_size = parse_positive_int(properties, "library.size", manifest_name) + chunk_size = parse_positive_int(properties, "chunk.size", manifest_name) + chunk_count = parse_positive_int(properties, "chunk.count", manifest_name) + expected_count = (library_size + chunk_size - 1) // chunk_size + if chunk_count != expected_count: + raise RuntimeError( + f"Invalid chunk.count in {manifest_name}: expected {expected_count}, " + f"found {chunk_count}" + ) + + output_path = os.path.join(output_dir, library_name) + output, temporary_path = temporary_output(output_dir, library_name) + total_bytes = 0 + try: + with output: + for index in range(chunk_count): + chunk_name = ( + f"{library_entry}{CHUNK_DIRECTORY_SUFFIX}/{index:05d}" + ) + expected_size = min(chunk_size, library_size - total_bytes) + crc_key = f"chunk.{index:05d}.crc32" + expected_crc_text = require_property(properties, crc_key, manifest_name) + try: + expected_crc = int(expected_crc_text, 16) + except ValueError as error: + raise RuntimeError( + f"Invalid {crc_key} in {manifest_name}: {expected_crc_text}" + ) from error + + actual_size = 0 + actual_crc = 0 + try: + source = archive.open(chunk_name) + except KeyError as error: + raise RuntimeError( + f"Missing native library chunk {chunk_name}" + ) from error + with source: + while True: + data = source.read(COPY_BUFFER_SIZE) + if not data: + break + output.write(data) + actual_size += len(data) + actual_crc = binascii.crc32(data, actual_crc) + + actual_crc &= 0xFFFFFFFF + if actual_size != expected_size: + raise RuntimeError( + f"Invalid size for {chunk_name}: expected {expected_size}, " + f"found {actual_size}" + ) + if actual_crc != expected_crc: + raise RuntimeError( + f"CRC32 mismatch for {chunk_name}: expected " + f"{expected_crc:08x}, found {actual_crc:08x}" + ) + total_bytes += actual_size + + if total_bytes != library_size: + raise RuntimeError( + f"Invalid reconstructed size for {library_name}: expected " + f"{library_size}, found {total_bytes}" + ) + install_output(temporary_path, output_path) + except Exception: + try: + os.remove(temporary_path) + except FileNotFoundError: + pass + raise + print( + f"Reconstructed {library_name} from {chunk_count} chunks -> {output_path}" + ) + + +def find_native_entries(archive): + conventional = {} + chunked = {} + for entry in archive.infolist(): + if entry.is_dir(): + continue + basename = os.path.basename(entry.filename) + if basename.endswith(MANIFEST_SUFFIX): + library_name = basename[: -len(MANIFEST_SUFFIX)] + if library_name == CUDF_LIBRARY_NAME: + chunked.setdefault(library_name, []).append(entry.filename) + elif basename == CUDF_LIBRARY_NAME: + conventional.setdefault(basename, []).append(entry.filename) + + duplicates = { + name: entries + for name, entries in {**conventional, **chunked}.items() + if len(entries) != 1 + } + if duplicates: + details = "; ".join( + f"{name}: {', '.join(entries)}" for name, entries in duplicates.items() + ) + raise RuntimeError(f"Multiple jar entries map to the same native library: {details}") + + overlap = sorted(set(conventional).intersection(chunked)) + if overlap: + raise RuntimeError( + "Jar contains conventional and chunked representations for: " + + ", ".join(overlap) + ) + return conventional, chunked + + +def extract_native_libraries(jar_path, output_dir): + os.makedirs(output_dir, exist_ok=True) + with zipfile.ZipFile(jar_path) as archive: + conventional, chunked = find_native_entries(archive) + if CUDF_LIBRARY_NAME not in conventional and CUDF_LIBRARY_NAME not in chunked: + raise RuntimeError( + "libcudf.so was not found as a conventional library or chunk manifest " + f"in {jar_path}" + ) + for library_name in sorted(conventional): + extract_conventional_library( + archive, conventional[library_name][0], output_dir + ) + for library_name in sorted(chunked): + reconstruct_chunked_library(archive, chunked[library_name][0], output_dir) + + +def main(): + parser = argparse.ArgumentParser() + parser.add_argument("jar", help="rapids-4-spark jar") + parser.add_argument("output_dir", help="directory for reconstructed native libraries") + args = parser.parse_args() + + try: + extract_native_libraries(args.jar, args.output_dir) + except (OSError, RuntimeError, UnicodeError, zipfile.BadZipFile) as error: + parser.exit(1, f"ERROR: {error}\n") + + +if __name__ == "__main__": + main() diff --git a/examples/UDF-Examples/RAPIDS-accelerated-UDFs/pom.xml b/examples/UDF-Examples/RAPIDS-accelerated-UDFs/pom.xml index 8b2f1648..0896466c 100644 --- a/examples/UDF-Examples/RAPIDS-accelerated-UDFs/pom.xml +++ b/examples/UDF-Examples/RAPIDS-accelerated-UDFs/pom.xml @@ -347,7 +347,6 @@ - @@ -387,24 +386,13 @@ - - - - - - - - - - - - - - - - - + + + + + +