From 41efa90a688311dc5dca9aadf9730820fce139a3 Mon Sep 17 00:00:00 2001 From: Julien Cuquemelle Date: Fri, 30 Jan 2026 18:21:58 +0100 Subject: [PATCH 1/2] Add 7zip support --- README.md | 1 + benchmark_zip_methods.py | 78 +++++++++++++++++++++++++++------------ cluster_pack/packaging.py | 62 ++++++++++++++++++++++++++++--- 3 files changed, 112 insertions(+), 29 deletions(-) diff --git a/README.md b/README.md index ff3d59e..599f044 100644 --- a/README.md +++ b/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 diff --git a/benchmark_zip_methods.py b/benchmark_zip_methods.py index 9817a6c..ef309f6 100755 --- a/benchmark_zip_methods.py +++ b/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,10 +10,16 @@ import shutil import subprocess import time -import zipfile 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_AVAILABLE, +) + PYTHON_VERSION = "3.11" TORCH_VERSION = "torch==2.8.0" BENCHMARK_DIR = "/tmp/zip_methods_benchmark" @@ -23,6 +29,8 @@ ("shutil", None), ("zipfile_cl0", 0), ("zipfile_cl1", 1), + ("7z_cl0", 0), + ("7z_cl6", 6), ] @@ -31,6 +39,9 @@ def check_uv_available() -> bool: return shutil.which("uv") is not None + + + def clear_cache(name: str = "pex") -> None: """Clear PEX/pip caches.""" pex_root = os.path.expanduser(f"~/.{name}") @@ -88,31 +99,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") +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). - 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 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 +138,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 +222,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_AVAILABLE: + 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") diff --git a/cluster_pack/packaging.py b/cluster_pack/packaging.py index d5aa7a1..ed7e618 100644 --- a/cluster_pack/packaging.py +++ b/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_AVAILABLE: bool = False class LayoutOptimizationParams(NamedTuple): @@ -145,7 +146,19 @@ def _detect_uv() -> bool: UV_AVAILABLE = _detect_uv() -def _make_zip_archive_zipfile(output_zip: str, source_dir: str, compress_level: int = 0) -> None: +def _detect_7zip() -> bool: + """Detect if 7z is installed and available in PATH.""" + if shutil.which("7z") is not None: + return True + else: + _logger.info("7z not found in PATH, falling back to single-threaded zip compression") + return False + + +SEVENZIP_AVAILABLE = _detect_7zip() + + +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 +173,59 @@ 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) + """ + cmd = [ + "7z", "a", + "-tzip", + f"-mx={compress_level}", + "-mmt=on", + output_zip, + os.path.join(source_dir, "*"), + ] + _logger.info(f"Creating zip archive with 7z (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_AVAILABLE: + 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: From fc880c0b1ae88d474b85d5b2a95a2a69014b3ce3 Mon Sep 17 00:00:00 2001 From: Julien Cuquemelle Date: Thu, 22 Jan 2026 11:26:35 +0100 Subject: [PATCH 2/2] Make 7zip detection work with 7z and 7za (depends on linux distro) And update CI to use 7z or not In order to not multiply test cases, we combine testing with uv and 7Zip, since both optims are orthogonal --- .github/workflows/main.yml | 2 ++ benchmark_zip_methods.py | 16 +++++++--------- cluster_pack/packaging.py | 32 ++++++++++++++++++-------------- 3 files changed, 27 insertions(+), 23 deletions(-) diff --git a/.github/workflows/main.yml b/.github/workflows/main.yml index 3fd714b..5a70df4 100644 --- a/.github/workflows/main.yml +++ b/.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 diff --git a/benchmark_zip_methods.py b/benchmark_zip_methods.py index ef309f6..760c502 100755 --- a/benchmark_zip_methods.py +++ b/benchmark_zip_methods.py @@ -10,6 +10,7 @@ import shutil import subprocess import time +from functools import partial from pathlib import Path from typing import List, Dict, Any @@ -17,20 +18,20 @@ make_zip_archive_shutil, make_zip_archive_zipfile, make_zip_archive_7z, - SEVENZIP_AVAILABLE, + 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), - ("7z_cl0", 0), - ("7z_cl6", 6), + ("shutil", None), ] @@ -39,9 +40,6 @@ def check_uv_available() -> bool: return shutil.which("uv") is not None - - - def clear_cache(name: str = "pex") -> None: """Clear PEX/pip caches.""" pex_root = os.path.expanduser(f"~/.{name}") @@ -222,7 +220,7 @@ 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_AVAILABLE: + if method_name.startswith("7z") and not SEVENZIP_EXECUTABLE: print(f"\n Skipping {method_name} (7z not available)") continue diff --git a/cluster_pack/packaging.py b/cluster_pack/packaging.py index ed7e618..3469ae1 100644 --- a/cluster_pack/packaging.py +++ b/cluster_pack/packaging.py @@ -50,7 +50,7 @@ class PythonEnvDescription(NamedTuple): LARGE_PEX_CMD = f"{UNPACKED_ENV_NAME}/__main__.py" UV_AVAILABLE: bool = False -SEVENZIP_AVAILABLE: bool = False +SEVENZIP_EXECUTABLE: Optional[str] = None class LayoutOptimizationParams(NamedTuple): @@ -143,19 +143,20 @@ def _detect_uv() -> bool: return False -UV_AVAILABLE = _detect_uv() - +def _detect_7zip() -> Optional[str]: + """Detect if 7z or 7za is installed and available in PATH. -def _detect_7zip() -> bool: - """Detect if 7z is installed and available in PATH.""" - if shutil.which("7z") is not None: - return True - else: - _logger.info("7z not found in PATH, falling back to single-threaded zip compression") - return False + :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 -SEVENZIP_AVAILABLE = _detect_7zip() +UV_AVAILABLE = _detect_uv() +SEVENZIP_EXECUTABLE = _detect_7zip() def make_zip_archive_zipfile(output_zip: str, source_dir: str, compress_level: int = 0) -> None: @@ -180,15 +181,17 @@ def make_zip_archive_7z(output_zip: str, source_dir: str, compress_level: int = :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 = [ - "7z", "a", + SEVENZIP_EXECUTABLE, "a", "-tzip", f"-mx={compress_level}", "-mmt=on", output_zip, os.path.join(source_dir, "*"), ] - _logger.info(f"Creating zip archive with 7z (compress_level={compress_level})") + _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: @@ -202,6 +205,7 @@ def make_zip_archive_shutil(output: str, source_dir: str) -> None: :param output: output path without .zip extension :param source_dir: directory to compress """ + shutil.make_archive(output, "zip", source_dir) @@ -218,7 +222,7 @@ def _make_zip_archive(output: str, source_dir: str, use_zipfile: bool, compress_ """ output_zip = output + ".zip" if use_zipfile: - if SEVENZIP_AVAILABLE: + 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})")