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/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..760c502 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,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") diff --git a/cluster_pack/packaging.py b/cluster_pack/packaging.py index d5aa7a1..3469ae1 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_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: