Skip to content

Chunk large native libraries for parallel extraction - #15356

Open
gerashegalov wants to merge 4 commits into
NVIDIA:mainfrom
gerashegalov:gerashegalov/chunk-native-libraries
Open

Chunk large native libraries for parallel extraction#15356
gerashegalov wants to merge 4 commits into
NVIDIA:mainfrom
gerashegalov:gerashegalov/chunk-native-libraries

Conversation

@gerashegalov

@gerashegalov gerashegalov commented Jul 23, 2026

Copy link
Copy Markdown
Collaborator

Description

Depends on rapidsai/cudf#23409, which teaches NativeDepsLoader to consume the generated format.

Related issue: #15145. This PR addresses its native JAR extraction bottleneck; background preloading and concurrent memory-pool initialization remain separate work.

Large CUDA native libraries are expensive to inflate through one sequential JAR stream. CUDA 13 also uses zstd-compressed fatbins, so additional whole-library DEFLATE compression is less valuable while startup latency remains significant.

This change makes chunking a cudf-spark distribution packaging concern:

  • native libraries at least 256 MiB are split into 32 MiB resources during process-resources
  • a versioned manifest records the library size, chunk geometry, and per-chunk CRC32 values
  • the existing create-parallel-worlds-jar execution picks up the generated resources and DEFLATEs them in its normal archive pass
  • the conventional resource is retained for smaller libraries and when JAR compression is disabled
  • a verify-phase task checks entry methods, resource presence, sizes, and reconstructed checksums

There is no new user-facing configuration. The POM properties are internal build controls so chunking remains limited to the cudf-spark distribution JAR.

Performance

Native extraction

Measured the exact pre-change and PR paths on the same host with JDK 17 and the same 1,504,051,352-byte CUDA12 libcudf.so. Each case ran six extractions in one JVM; the first invocation was treated as warm-up and the remaining five were used for the median.

Path First invocation Post-warm-up range Median
Pre-change conventional DEFLATED resource, 16 KiB sequential copy 6.411 s 6.370-6.426 s 6.371 s
45 DEFLATED chunks, 12-worker positional extraction 1.164 s 1.078-1.124 s 1.113 s

The post-warm-up median improves by 5.72x, reducing extraction latency by 82.5%. The first measured invocation improves by 5.51x. Both paths produced exactly 1,504,051,352 bytes.

JAR creation

Isolated the existing create-parallel-worlds-jar Maven execution using two resource trees on the same ext4 filesystem. All non-native files were hard-linked between the trees; only the native representation differed. Every invocation forced a new archive:

/usr/bin/time mvn jar:jar@create-parallel-worlds-jar -pl dist \
  -DskipTests -Dmaven.jar.forceCreation=true

Each case had one warm-up followed by five measured runs. Times below are direct-goal wall time, including identical Maven/JVM startup overhead.

Input representation Measured range Median Output JAR size
One conventional DEFLATED libcudf.so entry 54.87-55.14 s 55.07 s 952,504,790 bytes
45 DEFLATED chunks plus manifest 5.40-5.66 s 5.47 s 952,587,551 bytes

Chunking makes focused JAR creation 10.07x faster, reducing wall time by 90.1%. The chunked archive is only 82,761 bytes larger, a 0.0087% increase. Plexus Archiver can compress independent chunk entries concurrently instead of processing one 1.5 GB entry on a single worker.

Validation

  • ./build/make-scala-version-build-files.sh 2.13
  • mvn package -pl dist -DskipTests
  • mvn antrun:run@verify-native-library-chunks -pl dist -DskipTests -Dspark.rapids.source.basedir=$PWD
  • verifier reconstructed one native library from 45 DEFLATED chunks with matching checksums
  • exact cross-repository extraction benchmark described above
  • focused forced-rebuild JAR benchmark described above

Checklists

Documentation

  • Updated for new or modified user-facing features or behaviors
  • No user-facing change

Testing

  • Added or modified tests to cover new code paths
  • Covered by existing tests
  • Not required

Performance

  • Tests ran and results are added in the PR description
  • Issue filed with a link in the PR description
  • Not required

Signed-off-by: Gera Shegalov <gshegalov@nvidia.com>
@greptile-apps

greptile-apps Bot commented Jul 23, 2026

Copy link
Copy Markdown
Contributor

Greptile Summary

This PR splits large native libraries (≥ 256 MiB) into 32 MiB chunks during process-resources so that the existing create-parallel-worlds-jar step can DEFLATE each chunk concurrently. A versioned .chunks.properties manifest records size, geometry, and per-chunk CRC32 values; a verify-phase Ant/Jython task reconstructs and validates the output. The conventional (unchunked) resource is retained when libraries are small or when JAR compression is disabled.

  • chunk-native-libraries.py: Jython 2.7 script that walks the parallel-world directory two levels deep (arch → OS → files), splits qualifying .so/.dylib/.dll files into numbered chunks using Java FileInputStream/FileOutputStream, writes a .chunks.properties manifest with CRC32 per chunk, and atomically replaces the original library via rename/remove. Temporary directories and manifests are cleaned up on any failure path before the original exception is re-raised.
  • verify-native-library-chunks.py: Jython 2.7 script run in the verify phase that opens the produced JAR via ZipFile, checks DEFLATED entry presence and method, parses each manifest, and verifies per-chunk size and CRC32 before comparing the full verified-chunk set against the expected-chunk list.
  • dist/pom.xml / scala2.13/dist/pom.xml: Adds four internal POM properties, an initialize-phase metadata-dir deletion, a process-resources-phase chunk-native-libraries Ant/Jython execution, and a verify-phase verify-native-library-chunks execution; the latter is guarded by the presence of an enabled marker file so it is a no-op when no libraries were chunked.

Confidence Score: 5/5

Safe to merge. The chunking and verification logic is correct, error paths are clean, and the atomic rename/remove commit sequence is well-guarded by both the exception-handling structure and the initialize-phase parallel-world cleanup on the next invocation.

The chunking arithmetic is provably correct (chunk sizes sum exactly to library_size), CRC32 verification closes the integrity loop, the exception paths consistently use error-ignoring cleanup helpers so the original exception is never masked, and the enabled-marker guard ensures the verifier is a no-op when no chunking occurred. No logic defects were found in any of the four changed files.

Files Needing Attention: No files require special attention.

Important Files Changed

Filename Overview
dist/build/chunk-native-libraries.py New Jython 2.7 script that splits large native libraries into fixed-size chunks with per-chunk CRC32 and atomic rename/remove commit sequence; exception paths are clean (error-ignoring cleanup helper, outer finally for source stream).
dist/build/verify-native-library-chunks.py New Jython 2.7 verifier that opens the produced JAR, validates DEFLATED method, chunk geometry, per-chunk size, CRC32 integrity, and absence of the conventional unchunked resource; final set comparison catches any chunk-count drift between the metadata list files and manifests.
dist/pom.xml Adds four internal chunking POM properties, metadata-dir cleanup in the initialize phase, a process-resources Jython execution for chunking (guarded by both dist.native.chunking.enabled and dist.jar.compress), and a verify-phase execution guarded by the enabled marker; all wired into the existing Ant lifecycle cleanly.
scala2.13/dist/pom.xml Identical diff to dist/pom.xml applied to the Scala 2.13 mirror POM; correctly keeps the two POMs in sync.

Sequence Diagram

sequenceDiagram
    participant Maven
    participant Init as initialize phase
    participant GenRes as generate-resources
    participant ProcRes as process-resources
    participant Jar as package phase
    participant Verify as verify phase

    Maven->>Init: clean-any-prior-output
    Init->>Init: delete parallel-world, jars, deps, native-chunks
    Maven->>GenRes: create-parallel-world
    GenRes->>GenRes: populate parallel-world from aggregator JARs
    Maven->>ProcRes: copy-jni-and-ucx-classes
    ProcRes->>ProcRes: copy jni-deps into parallel-world
    Maven->>ProcRes: chunk-native-libraries (Jython, guarded by enabled+compress)
    ProcRes->>ProcRes: walk parallel-world arch/os files
    ProcRes->>ProcRes: "split each library >=256MiB into 32MiB chunks"
    ProcRes->>ProcRes: write manifest with library size and per-chunk CRC32
    ProcRes->>ProcRes: atomic rename tmp dirs, remove original library
    ProcRes->>ProcRes: write deflated-chunks.list and chunk-manifests.list
    Maven->>Jar: create-parallel-worlds-jar
    Jar->>Jar: DEFLATE all parallel-world entries including chunks
    Maven->>Verify: verify-native-library-chunks (Jython, guarded by enabled marker)
    Verify->>Verify: open dist JAR and check all chunks are DEFLATED
    Verify->>Verify: parse each manifest and verify chunk count
    Verify->>Verify: verify no conventional resource beside manifest
    Verify->>Verify: read each chunk and verify size and CRC32
    Verify->>Verify: compare verified-chunk set against expected list
Loading

Reviews (4): Last reviewed commit: "Support byte-size suffixes for native ch..." | Re-trigger Greptile

Comment thread dist/build/chunk-native-libraries.py
Comment thread dist/build/verify-native-library-chunks.py
Comment thread dist/build/chunk-native-libraries.py
@sameerz sameerz added the performance A performance related task/issue label Jul 23, 2026
Signed-off-by: Gera Shegalov <gshegalov@nvidia.com>
@gerashegalov
gerashegalov requested a review from a team July 23, 2026 16:35
@gerashegalov gerashegalov self-assigned this Jul 23, 2026
Signed-off-by: Gera Shegalov <gshegalov@nvidia.com>
@nvauto

nvauto commented Jul 27, 2026

Copy link
Copy Markdown
Collaborator

NOTE: release/26.08 has been created from main. Please retarget your PR to release/26.08 if it should be included in the release.

Comment thread scala2.13/dist/pom.xml Outdated
Comment thread dist/build/verify-native-library-chunks.py
Signed-off-by: Gera Shegalov <gshegalov@nvidia.com>
rapids-bot Bot pushed a commit to rapidsai/cudf that referenced this pull request Jul 30, 2026
Add an internal chunk-manifest resource format for large native libraries and extract those chunks concurrently. The loader:

- preserves conventional single-resource loading and `ai.rapids.cudf.lib-native-dir`
- reads up to 12 JAR entries concurrently
- writes directly to the pre-sized output with positional `FileChannel` writes
- validates manifest structure, exact chunk sizes, and per-chunk CRC32 values
- removes partial output after extraction failures

Related issue: [NVIDIA/cudf-spark#15145](NVIDIA/cudf-spark#15145) identifies synchronous, single-threaded native JAR extraction as a major executor startup cost.

This allows distribution JARs containing 900+ MiB CUDA native libraries to trade a measured 0.0087% increase in archive size for substantially faster startup and JAR creation. NVIDIA/cudf-spark#15356 is the companion producer PR.

### Performance

#### Native extraction

Measured the exact pre-change and PR paths on the same host with JDK 17 and the same 1,504,051,352-byte CUDA12 `libcudf.so`. Each case ran six extractions in one JVM; the first invocation was treated as warm-up and the remaining five were used for the median.

| Path | First invocation | Post-warm-up range | Median |
| --- | ---: | ---: | ---: |
| Pre-change conventional DEFLATED resource, 16 KiB sequential copy | 6.411 s | 6.370-6.426 s | 6.371 s |
| 45 DEFLATED chunks, 12-worker positional extraction | 1.164 s | 1.078-1.124 s | 1.113 s |

The post-warm-up median improves by **5.72x**, reducing extraction latency by **82.5%**. The first measured invocation improves by 5.51x. Both paths produced exactly 1,504,051,352 bytes.

#### Archive size and creation time

Using matched resource trees and forcing the existing `jar:jar@create-parallel-worlds-jar` execution:

| Input representation | Median | Output JAR size |
| --- | ---: | ---: |
| One conventional DEFLATED `libcudf.so` entry | 55.07 s | 952,504,790 bytes |
| 45 DEFLATED chunks plus manifest | 5.47 s | 952,587,551 bytes |

Chunking increased the archive by 82,761 bytes (**0.0087%**) while making focused JAR creation **10.07x faster**, a **90.1%** wall-time reduction. The full benchmark methodology and measured ranges are in NVIDIA/cudf-spark#15356.

### Validation

- `mvn compiler:compile compiler:testCompile`
- `mvn surefire:test -Dtest=NativeDepsLoaderExtractionTest` (10 tests)
- `mvn surefire:test@native-deps-loader-test` (5 tests)
- Cross-repository production JAR benchmark described above

The full native Maven lifecycle was not run locally because the installed CMake is 3.28 and this branch requires CMake 4.0; the Java sources and focused tests were compiled and run directly.

Authors:
  - Gera Shegalov (https://github.com/gerashegalov)

Approvers:
  - Zach Puller (https://github.com/zpuller)
  - Nghia Truong (https://github.com/ttnghia)

URL: #23409

@abellina abellina left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

LGTM

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

performance A performance related task/issue

Projects

None yet

Development

Successfully merging this pull request may close these issues.

5 participants