Add 7zip support

.github/workflows/main.yml

@@ -60,11 +60,13 @@ jobs:
         uv pip install --system -e .
         uv pip install --system -r tests-requirements.txt
     - name: Install dependencies (without uv)
+      # Also remove 7z to test fallback path
       if: ${{ !matrix.use-uv }}
       run: |
         pip install --upgrade pip
         pip install -e .
         pip install -r tests-requirements.txt
+        sudo mv /usr/bin/7z /usr/bin/7z.bak
     - name: Tests
       run: |
         pytest -m "not hadoop and not conda" -s tests

README.md

@@ -32,6 +32,7 @@ $ pip install .
 
 - Cluster-pack supports Python ≥3.9.
 - Cluster-pack can speed up pex creation by using uv if available.
+- Cluster-pack can use 7z for faster multithreaded zip compression if available.
 
 ## Feature flags
 - C_PACK_USER: override the current user for HDFS path generation and Skein impersonation

benchmark_zip_methods.py

@@ -1,5 +1,5 @@
 #!/usr/bin/env python3
-"""Benchmark pex creation comparing shutil vs zipfile for zipping, across 3 pex layouts.
+"""Benchmark pex creation comparing shutil vs zipfile vs 7z for zipping, across 3 pex layouts.
 
 This benchmark only runs when uv is available and uses --max-install-jobs 0 setting.
 
@@ -10,19 +10,28 @@
 import shutil
 import subprocess
 import time
-import zipfile
+from functools import partial
 from pathlib import Path
 from typing import List, Dict, Any
 
+from cluster_pack.packaging import (
+    make_zip_archive_shutil,
+    make_zip_archive_zipfile,
+    make_zip_archive_7z,
+    SEVENZIP_EXECUTABLE,
+)
+
 PYTHON_VERSION = "3.11"
 TORCH_VERSION = "torch==2.8.0"
 BENCHMARK_DIR = "/tmp/zip_methods_benchmark"
 
-PEX_LAYOUTS = ["zipapp", "packed", "loose"]
+PEX_LAYOUTS = ["packed", "loose", "zipapp",]
 ZIP_METHODS = [
-    ("shutil", None),
+    ("7z_cl0", 0),
+    ("7z_cl1", 1),
     ("zipfile_cl0", 0),
     ("zipfile_cl1", 1),
+    ("shutil", None),
 ]
 
 
@@ -88,31 +97,23 @@ def create_pex_with_layout(venv_path: str, output_path: str, layout: str) -> flo
     return elapsed
 
 
-def zip_with_shutil(source_dir: str, output: str) -> tuple[float, float]:
-    """Zip using shutil.make_archive. Returns (time, size_mb)."""
-    if os.path.exists(output + ".zip"):
-        os.remove(output + ".zip")
-
-    start = time.time()
-    shutil.make_archive(output, "zip", source_dir)
-    elapsed = time.time() - start
-    size_mb = os.path.getsize(output + ".zip") / (1024 * 1024)
-    return elapsed, size_mb
+def timed_zip(zip_func, source_dir: str, output: str, compress_level: int = None) -> tuple[float, float]:
+    """Time a zip function and return (elapsed_time, size_mb).
 
-
-def zip_with_zipfile(source_dir: str, output: str, compresslevel: int = 1) -> tuple[float, float]:
-    """Zip using zipfile module. Returns (time, size_mb)."""
+    :param zip_func: one of make_zip_archive_shutil, make_zip_archive_zipfile, make_zip_archive_7z
+    :param source_dir: directory to compress
+    :param output: output path without .zip extension
+    :param compress_level: compression level (only for zipfile/7z)
+    """
     output_zip = output + ".zip"
     if os.path.exists(output_zip):
         os.remove(output_zip)
 
     start = time.time()
-    with zipfile.ZipFile(output_zip, 'w', zipfile.ZIP_DEFLATED, compresslevel=compresslevel) as zf:
-        for root, _, files in os.walk(source_dir):
-            for f in files:
-                full_path = os.path.join(root, f)
-                arc_name = os.path.relpath(full_path, source_dir)
-                zf.write(full_path, arc_name)
+    if zip_func == make_zip_archive_shutil:
+        zip_func(output, source_dir)
+    else:
+        zip_func(output_zip, source_dir, compress_level=compress_level)
     elapsed = time.time() - start
     size_mb = os.path.getsize(output_zip) / (1024 * 1024)
     return elapsed, size_mb
@@ -135,6 +136,23 @@ def unzip_archive(zip_path: str, extract_dir: str) -> float:
     return elapsed
 
 
+def unzip_with_7z(zip_path: str, extract_dir: str) -> float:
+    """Unzip using 7z. Returns time in seconds."""
+    if os.path.exists(extract_dir):
+        shutil.rmtree(extract_dir)
+    os.makedirs(extract_dir)
+
+    start = time.time()
+    subprocess.run(
+        ["7z", "x", "-y", "-mmt=on", f"-o{extract_dir}", zip_path],
+        check=True,
+        stdout=subprocess.DEVNULL,
+        stderr=subprocess.DEVNULL,
+    )
+    elapsed = time.time() - start
+    return elapsed
+
+
 def get_dir_size_mb(path: str) -> float:
     """Get directory size in MB."""
     return sum(f.stat().st_size for f in Path(path).rglob('*') if f.is_file()) / (1024 * 1024)
@@ -202,18 +220,28 @@ def benchmark_layout(venv_path: str, layout: str) -> List[Dict[str, Any]]:
         print(f"  Uncompressed size: {uncompressed_size_mb:.1f}MB")
 
         for method_name, compresslevel in ZIP_METHODS:
+            if method_name.startswith("7z") and not SEVENZIP_EXECUTABLE:
+                print(f"\n  Skipping {method_name} (7z not available)")
+                continue
+
             zip_output = os.path.join(BENCHMARK_DIR, f"{layout}_{method_name}")
 
             print(f"\n  Testing {method_name}...")
             if method_name == "shutil":
-                compress_time, compressed_size_mb = zip_with_shutil(output_path, zip_output)
+                zip_func = make_zip_archive_shutil
+            elif method_name.startswith("7z"):
+                zip_func = make_zip_archive_7z
             else:
-                compress_time, compressed_size_mb = zip_with_zipfile(output_path, zip_output, compresslevel)
+                zip_func = make_zip_archive_zipfile
+            compress_time, compressed_size_mb = timed_zip(zip_func, output_path, zip_output, compresslevel)
 
             ratio = (compressed_size_mb / uncompressed_size_mb) * 100
 
             extract_dir = os.path.join(BENCHMARK_DIR, f"{layout}_{method_name}_extract")
-            extract_time = unzip_archive(zip_output + ".zip", extract_dir)
+            if method_name.startswith("7z"):
+                extract_time = unzip_with_7z(zip_output + ".zip", extract_dir)
+            else:
+                extract_time = unzip_archive(zip_output + ".zip", extract_dir)
 
             print(f"    Compress: {compress_time:.2f}s, Size: {compressed_size_mb:.1f}MB ({ratio:.1f}%), Extract: {extract_time:.2f}s")
 

cluster_pack/packaging.py

@@ -50,6 +50,7 @@ class PythonEnvDescription(NamedTuple):
 LARGE_PEX_CMD = f"{UNPACKED_ENV_NAME}/__main__.py"
 
 UV_AVAILABLE: bool = False
+SEVENZIP_EXECUTABLE: Optional[str] = None
 
 
 class LayoutOptimizationParams(NamedTuple):
@@ -142,10 +143,23 @@ def _detect_uv() -> bool:
         return False
 
 
+def _detect_7zip() -> Optional[str]:
+    """Detect if 7z or 7za is installed and available in PATH.
+
+    :return: the executable name ('7z' or '7za') if found, None otherwise
+    """
+    for executable in ["7z", "7za"]:
+        if shutil.which(executable) is not None:
+            return executable
+    _logger.info("7z/7za not found in PATH, falling back to single-threaded zip compression")
+    return None
+
+
 UV_AVAILABLE = _detect_uv()
+SEVENZIP_EXECUTABLE = _detect_7zip()
 
 
-def _make_zip_archive_zipfile(output_zip: str, source_dir: str, compress_level: int = 0) -> None:
+def make_zip_archive_zipfile(output_zip: str, source_dir: str, compress_level: int = 0) -> None:
     """Create a zip archive using zipfile module with specified compression level.
 
     :param output_zip: output path (with .zip extension)
@@ -160,20 +174,62 @@ def _make_zip_archive_zipfile(output_zip: str, source_dir: str, compress_level:
                 zf.write(full_path, arc_name)
 
 
+def make_zip_archive_7z(output_zip: str, source_dir: str, compress_level: int = 6) -> None:
+    """Create a zip archive using 7z with multithreading.
+
+    :param output_zip: output path (with .zip extension)
+    :param source_dir: directory to compress
+    :param compress_level: compression level 0-9 (0=store, 1=fastest, 9=best)
+    """
+    if not SEVENZIP_EXECUTABLE:
+        raise RuntimeError("7z executable not found")
+    cmd = [
+        SEVENZIP_EXECUTABLE, "a",
+        "-tzip",
+        f"-mx={compress_level}",
+        "-mmt=on",
+        output_zip,
+        os.path.join(source_dir, "*"),
+    ]
+    _logger.info(f"Creating zip archive with {SEVENZIP_EXECUTABLE} (compress_level={compress_level})")
+    try:
+        subprocess.run(cmd, cwd=source_dir, stderr=subprocess.PIPE, stdout=subprocess.PIPE, check=True)
+    except subprocess.CalledProcessError as e:
+        _logger.warning(f"7z failed: {e.stderr.decode('utf-8', errors='replace')}, falling back to zipfile")
+        make_zip_archive_zipfile(output_zip, source_dir, compress_level=compress_level)
+
+
+def make_zip_archive_shutil(output: str, source_dir: str) -> None:
+    """Create a zip archive using shutil.make_archive.
+
+    :param output: output path without .zip extension
+    :param source_dir: directory to compress
+    """
+
+    shutil.make_archive(output, "zip", source_dir)
+
+
 def _make_zip_archive(output: str, source_dir: str, use_zipfile: bool, compress_level: int) -> None:
     """Create a zip archive from source_dir.
 
+    Uses 7z with multithreading if available and use_zipfile=True, otherwise falls back to
+    zipfile module or shutil.make_archive.
+
     :param output: output path without .zip extension
     :param source_dir: directory to compress
-    :param use_zipfile: True to use zipfile module, False for shutil.make_archive
+    :param use_zipfile: True to use zipfile/7z, False for shutil.make_archive
     :param compress_level: compression level (0=store, 1-9=compression), only used if use_zipfile=True
     """
+    output_zip = output + ".zip"
     if use_zipfile:
-        _logger.info(f"Creating zip archive with zipfile (compresslevel={compress_level})")
-        _make_zip_archive_zipfile(output + ".zip", source_dir, compress_level=compress_level)
+        if SEVENZIP_EXECUTABLE:
+            make_zip_archive_7z(output_zip, source_dir, compress_level=compress_level)
+        else:
+            _logger.info(f"Creating zip archive with zipfile (compresslevel={compress_level})")
+            make_zip_archive_zipfile(output_zip, source_dir, compress_level=compress_level)
     else:
         _logger.info("Creating zip archive with shutil.make_archive")
-        shutil.make_archive(output, "zip", source_dir)
+        make_zip_archive_shutil(output, source_dir)
 
 
 def _get_tmp_dir() -> str: