Chunk large native libraries for parallel extraction

dist/build/chunk-native-libraries.py

@@ -0,0 +1,207 @@
+# 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.
+
+import os
+import re
+import shutil
+
+from java.io import FileInputStream, FileOutputStream
+from java.util.zip import CRC32
+from jarray import zeros
+
+
+COPY_BUFFER_SIZE = 1024 * 1024
+MANIFEST_SUFFIX = ".chunks.properties"
+CHUNK_DIRECTORY_SUFFIX = ".chunks"
+NATIVE_SUFFIXES = (".so", ".dylib", ".dll")
+EMPTY_INCLUDE_PATTERN = "__no_native_chunks__"
+BYTE_SIZE_RE = re.compile(r"^\s*(\d+)\s*([KMGT]?)B?\s*$", re.IGNORECASE)
+BYTE_SIZE_MULTIPLIERS = {
+    "": 1,
+    "K": 1024,
+    "M": 1024 ** 2,
+    "G": 1024 ** 3,
+    "T": 1024 ** 4,
+}
+
+
+def parse_byte_size(value):
+    match = BYTE_SIZE_RE.match(str(value))
+    if not match:
+        raise ValueError("Invalid byte size: %r" % value)
+    amount, suffix = match.groups()
+    return long(amount) * BYTE_SIZE_MULTIPLIERS[suffix.upper()]
+
+
+def ensure_directory(path):
+    if not os.path.isdir(path):
+        os.makedirs(path)
+
+
+def remove_path(path):
+    if os.path.isdir(path):
+        shutil.rmtree(path)
+    elif os.path.exists(path):
+        os.remove(path)
+
+
+def remove_path_ignoring_errors(path):
+    try:
+        remove_path(path)
+    except:
+        pass
+
+
+def native_libraries(root_dir, minimum_size):
+    candidates = []
+    for arch in sorted(os.listdir(root_dir)):
+        arch_dir = os.path.join(root_dir, arch)
+        if not os.path.isdir(arch_dir):
+            continue
+        for operating_system in sorted(os.listdir(arch_dir)):
+            os_dir = os.path.join(arch_dir, operating_system)
+            if not os.path.isdir(os_dir):
+                continue
+            for name in sorted(os.listdir(os_dir)):
+                path = os.path.join(os_dir, name)
+                if (os.path.isfile(path) and name.endswith(NATIVE_SUFFIXES)
+                        and os.path.getsize(path) >= minimum_size):
+                    candidates.append(path)
+    return candidates
+
+
+def split_library(root_dir, library_path, chunk_size):
+    relative_library = os.path.relpath(library_path, root_dir).replace(os.sep, "/")
+    source_stat = os.stat(library_path)
+    library_size = source_stat.st_size
+    chunk_dir = library_path + CHUNK_DIRECTORY_SUFFIX
+    temporary_chunk_dir = chunk_dir + ".tmp"
+    manifest_path = library_path + MANIFEST_SUFFIX
+    temporary_manifest = manifest_path + ".tmp"
+
+    remove_path(temporary_chunk_dir)
+    remove_path(temporary_manifest)
+    if os.path.exists(chunk_dir) or os.path.exists(manifest_path):
+        raise RuntimeError("Chunk output already exists for %s" % relative_library)
+    ensure_directory(temporary_chunk_dir)
+
+    source = FileInputStream(library_path)
+    buffer = zeros(COPY_BUFFER_SIZE, "b")
+    deflated_entries = []
+    chunk_crc32 = []
+    chunk_count = 0
+    total_bytes = 0
+    try:
+        while total_bytes < library_size:
+            chunk_name = "%05d" % chunk_count
+            chunk_path = os.path.join(temporary_chunk_dir, chunk_name)
+            chunk_output = FileOutputStream(chunk_path)
+            chunk_crc = CRC32()
+            chunk_bytes = 0
+            expected = min(chunk_size, library_size - total_bytes)
+            try:
+                while chunk_bytes < expected:
+                    requested = int(min(COPY_BUFFER_SIZE, expected - chunk_bytes))
+                    count = source.read(buffer, 0, requested)
+                    if count < 0:
+                        raise RuntimeError("Unexpected end of native library %s" % relative_library)
+                    if count:
+                        chunk_output.write(buffer, 0, count)
+                        chunk_crc.update(buffer, 0, count)
+                        chunk_bytes += count
+                        total_bytes += count
+            finally:
+                chunk_output.close()
+
+            os.utime(chunk_path, (source_stat.st_atime, source_stat.st_mtime))
+            final_relative = (
+                relative_library + CHUNK_DIRECTORY_SUFFIX + "/" + chunk_name)
+            deflated_entries.append(final_relative)
+            chunk_crc32.append(chunk_crc.getValue())
+            chunk_count += 1
+    except:
+        remove_path_ignoring_errors(temporary_chunk_dir)
+        remove_path_ignoring_errors(temporary_manifest)
+        raise
+    finally:
+        source.close()
+
+    if total_bytes != library_size:
+        remove_path(temporary_chunk_dir)
+        raise RuntimeError(
+            "Native library changed while chunking %s: expected %d bytes, read %d"
+            % (relative_library, library_size, total_bytes))
+
+    manifest = (
+        "format.version=1\n"
+        "library.size=%d\n"
+        "chunk.size=%d\n"
+        "chunk.count=%d\n"
+        % (library_size, chunk_size, chunk_count))
+    manifest += "".join(
+        "chunk.%05d.crc32=%08x\n" % (index, value)
+        for index, value in enumerate(chunk_crc32))
+    manifest_output = open(temporary_manifest, "w")
+    try:
+        manifest_output.write(manifest)
+    finally:
+        manifest_output.close()
+    os.utime(temporary_manifest, (source_stat.st_atime, source_stat.st_mtime))
+
+    # Any failure here aborts the Maven execution before JAR creation. The initialize-phase
+    # cleanup removes partial parallel-world output before the next invocation.
+    os.rename(temporary_chunk_dir, chunk_dir)
+    os.rename(temporary_manifest, manifest_path)
+    os.remove(library_path)
+    return relative_library, deflated_entries
+
+
+def write_lines(path, values):
+    output = open(path, "w")
+    try:
+        if not values:
+            output.write(EMPTY_INCLUDE_PATTERN)
+            output.write("\n")
+        for value in sorted(values):
+            output.write(value)
+            output.write("\n")
+    finally:
+        output.close()
+
+
+root_dir = attributes.get("root_dir")
+metadata_dir = attributes.get("metadata_dir")
+minimum_size = parse_byte_size(attributes.get("minimum_size"))
+chunk_size = parse_byte_size(attributes.get("chunk_size"))
+
+if minimum_size <= 0 or chunk_size <= 0:
+    raise RuntimeError("Native chunk sizes must be positive")
+
+ensure_directory(metadata_dir)
+deflated_entries = []
+manifests = []
+for library_path in native_libraries(root_dir, minimum_size):
+    relative_library, library_deflated = split_library(
+        root_dir, library_path, chunk_size)
+    deflated_entries.extend(library_deflated)
+    manifests.append(relative_library + MANIFEST_SUFFIX)
+    self.log(
+        "Chunked %s into %d DEFLATED entries"
+        % (relative_library, len(library_deflated)))
+
+write_lines(os.path.join(metadata_dir, "deflated-chunks.list"), deflated_entries)
+write_lines(os.path.join(metadata_dir, "chunk-manifests.list"), manifests)
+if manifests:
+    marker = open(os.path.join(metadata_dir, "enabled"), "w")
+    marker.close()

dist/build/verify-native-library-chunks.py

@@ -0,0 +1,142 @@
+# 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.
+
+import os
+
+from java.util import Properties
+from java.util.zip import CRC32, ZipEntry, ZipFile
+from jarray import zeros
+
+
+COPY_BUFFER_SIZE = 1024 * 1024
+MANIFEST_SUFFIX = ".chunks.properties"
+EMPTY_INCLUDE_PATTERN = "__no_native_chunks__"
+
+
+def read_lines(path):
+    source = open(path, "r")
+    try:
+        return set(
+            line.strip() for line in source
+            if line.strip() and line.strip() != EMPTY_INCLUDE_PATTERN)
+    finally:
+        source.close()
+
+
+def require_property(properties, key, manifest_name):
+    value = properties.getProperty(key)
+    if value is None or not value.strip():
+        raise RuntimeError("Missing %s in %s" % (key, manifest_name))
+    return value.strip()
+
+
+jar_path = attributes.get("jar_path")
+metadata_dir = attributes.get("metadata_dir")
+deflated_entries = read_lines(os.path.join(metadata_dir, "deflated-chunks.list"))
+manifest_entries = read_lines(os.path.join(metadata_dir, "chunk-manifests.list"))
+expected_chunks = deflated_entries
+verified_chunks = set()
+
+archive = ZipFile(jar_path)
+buffer = zeros(COPY_BUFFER_SIZE, "b")
+try:
+    for entry_name in sorted(deflated_entries):
+        entry = archive.getEntry(entry_name)
+        if entry is None:
+            raise RuntimeError("Missing DEFLATED native chunk %s" % entry_name)
+        if entry.getMethod() != ZipEntry.DEFLATED:
+            raise RuntimeError("Native chunk %s is not DEFLATED" % entry_name)
+
+    for manifest_name in sorted(manifest_entries):
+        manifest_entry = archive.getEntry(manifest_name)
+        if manifest_entry is None:
+            raise RuntimeError("Missing native chunk manifest %s" % manifest_name)
+        properties = Properties()
+        manifest_input = archive.getInputStream(manifest_entry)
+        try:
+            properties.load(manifest_input)
+        finally:
+            manifest_input.close()
+
+        if require_property(properties, "format.version", manifest_name) != "1":
+            raise RuntimeError("Unsupported chunk manifest version in %s" % manifest_name)
+        library_size = long(require_property(properties, "library.size", manifest_name))
+        chunk_size = long(require_property(properties, "chunk.size", manifest_name))
+        chunk_count = int(require_property(properties, "chunk.count", manifest_name))
+        if library_size <= 0:
+            raise RuntimeError("Invalid library.size in %s" % manifest_name)
+        if chunk_size <= 0:
+            raise RuntimeError("Invalid chunk.size in %s" % manifest_name)
+        if chunk_count <= 0:
+            raise RuntimeError("Invalid chunk.count in %s" % manifest_name)
+        expected_count = 1 + ((library_size - 1) // chunk_size)
+        if chunk_count != expected_count:
+            raise RuntimeError("Invalid chunk count in %s" % manifest_name)
+
+        library_name = manifest_name[:-len(MANIFEST_SUFFIX)]
+        if archive.getEntry(library_name) is not None:
+            raise RuntimeError(
+                "Conventional native resource still exists beside %s" % manifest_name)
+
+        total_size = 0
+        for index in range(chunk_count):
+            chunk_name = "%s.chunks/%05d" % (library_name, index)
+            chunk_crc_key = "chunk.%05d.crc32" % index
+            expected_chunk_crc = long(
+                require_property(properties, chunk_crc_key, manifest_name), 16)
+            chunk_entry = archive.getEntry(chunk_name)
+            if chunk_entry is None:
+                raise RuntimeError("Missing native chunk %s" % chunk_name)
+            expected_size = (
+                chunk_size if index < chunk_count - 1
+                else library_size - chunk_size * index)
+            if chunk_entry.getSize() != expected_size:
+                raise RuntimeError(
+                    "Native chunk %s has size %d, expected %d"
+                    % (chunk_name, chunk_entry.getSize(), expected_size))
+            chunk_input = archive.getInputStream(chunk_entry)
+            chunk_bytes = 0
+            chunk_crc = CRC32()
+            try:
+                while True:
+                    count = chunk_input.read(buffer)
+                    if count < 0:
+                        break
+                    if count:
+                        chunk_crc.update(buffer, 0, count)
+                        chunk_bytes += count
+            finally:
+                chunk_input.close()
+            if chunk_bytes != expected_size:
+                raise RuntimeError(
+                    "Native chunk %s read %d bytes, expected %d"
+                    % (chunk_name, chunk_bytes, expected_size))
+            if chunk_crc.getValue() != expected_chunk_crc:
+                raise RuntimeError("Native chunk CRC mismatch for %s" % chunk_name)
+            total_size += chunk_bytes
+            verified_chunks.add(chunk_name)
+
+        if total_size != library_size:
+            raise RuntimeError(
+                "Reconstructed native library size mismatch for %s" % library_name)
+finally:
+    archive.close()
+
+if verified_chunks != expected_chunks:
+    raise RuntimeError(
+        "Native chunk lists do not match manifests: expected=%d verified=%d"
+        % (len(expected_chunks), len(verified_chunks)))
+self.log(
+    "Verified %d chunked native libraries and %d chunks"
+    % (len(manifest_entries), len(verified_chunks)))