Skip to content

Commit

Permalink
parallel get_wheel_elfdata
Browse files Browse the repository at this point in the history
  • Loading branch information
oraluben committed Mar 7, 2025
1 parent b586739 commit 7d8abfd
Show file tree
Hide file tree
Showing 2 changed files with 63 additions and 61 deletions.
26 changes: 7 additions & 19 deletions src/auditwheel/pool.py
Original file line number Diff line number Diff line change
@@ -1,4 +1,3 @@
import time
from concurrent.futures import Future, ThreadPoolExecutor
from pathlib import Path
from typing import Any, Callable, Optional
Expand All @@ -21,16 +20,16 @@ class FileTaskExecutor:
>>> executor.wait() # Wait for all tasks to complete
"""

def __init__(self, concurrent: int = 1):
def __init__(self, concurrent: int = 0):
self.executor = (
None
if concurrent == 1
else ThreadPoolExecutor(concurrent if concurrent > 1 else None)
)
self.working_map: dict[Path, Future[tuple[str, str]]] = {}
self.working_map: dict[Path, Future[Any]] = {}

def submit(
self, path: Path, fn: Callable[[Any], Any], /, *args: Any, **kwargs: Any
self, path: Path, fn: Callable[..., Any], /, *args: Any, **kwargs: Any
) -> None:
if not path.is_absolute():
path = path.absolute()
Expand All @@ -39,12 +38,11 @@ def submit(
if self.executor is None:
future = Future()
future.set_result(fn(*args, **kwargs))
return None
return

assert path not in self.working_map, "path already in working_map"
future = self.executor.submit(fn, *args, **kwargs)
self.working_map[path] = future
return future

def wait(self, path: Optional[Path] = None) -> None:
"""Wait for tasks to complete.
Expand All @@ -62,18 +60,8 @@ def wait(self, path: Optional[Path] = None) -> None:
for future in self.working_map.values():
future.result()
self.working_map.clear()
elif future := self.working_map.pop(path, None):
future.result()


def fake_job(i: int) -> int:
print(f"start {i}")
time.sleep(i)
print(f"end {i}")
elif path in self.working_map:
self.working_map.pop(path).result()


if __name__ == "__main__":
executor = FileTaskExecutor(concurrent=0)
for i in range(10):
executor.submit(Path(f"test{i}.txt"), fake_job, i)
executor.wait()
POOL = FileTaskExecutor()
98 changes: 56 additions & 42 deletions src/auditwheel/wheel_abi.py
Original file line number Diff line number Diff line change
Expand Up @@ -10,6 +10,10 @@
from pathlib import Path
from typing import Optional, TypeVar

from elftools.elf.elffile import ELFFile

from auditwheel.pool import POOL

from . import json
from .architecture import Architecture
from .elfutils import (
Expand Down Expand Up @@ -94,6 +98,55 @@ def get_wheel_elfdata(
shared_libraries_with_invalid_machine = []

platform_wheel = False

def inner(fn: Path, elf: ELFFile) -> None:
nonlocal \
platform_wheel, \
shared_libraries_in_purelib, \
uses_ucs2_symbols, \
uses_PyFPE_jbuf

so_name = fn.name

log.debug("processing: %s", fn)
elftree = ldd(fn, exclude=exclude)

try:
arch = elftree.platform.baseline_architecture
if arch != wheel_policy.architecture.baseline:
shared_libraries_with_invalid_machine.append(so_name)
log.warning("ignoring: %s with %s architecture", so_name, arch)
return
except ValueError:
shared_libraries_with_invalid_machine.append(so_name)
log.warning("ignoring: %s with unknown architecture", so_name)
return

platform_wheel = True

for key, value in elf_find_versioned_symbols(elf):
log.debug("key %s, value %s", key, value)
versioned_symbols[key].add(value)

is_py_ext, py_ver = elf_is_python_extension(fn, elf)

# If the ELF is a Python extention, we definitely need to
# include its external dependencies.
if is_py_ext:
full_elftree[fn] = elftree
uses_PyFPE_jbuf |= elf_references_PyFPE_jbuf(elf)
if py_ver == 2:
uses_ucs2_symbols |= any(True for _ in elf_find_ucs2_symbols(elf))
full_external_refs[fn] = wheel_policy.lddtree_external_references(
elftree, ctx.path
)
else:
# If the ELF is not a Python extension, it might be
# included in the wheel already because auditwheel repair
# vendored it, so we will check whether we should include
# its internal references later.
nonpy_elftree[fn] = elftree

for fn, elf in elf_file_filter(ctx.iter_files()):
# Check for invalid binary wheel format: no shared library should
# be found in purelib
Expand All @@ -104,49 +157,8 @@ def get_wheel_elfdata(
if any(p.name == "purelib" for p in fn.parents):
shared_libraries_in_purelib.append(so_name)

# If at least one shared library exists in purelib, this is going
# to fail and there's no need to do further checks
if not shared_libraries_in_purelib:
log.debug("processing: %s", fn)
elftree = ldd(fn, exclude=exclude)

try:
arch = elftree.platform.baseline_architecture
if arch != wheel_policy.architecture.baseline:
shared_libraries_with_invalid_machine.append(so_name)
log.warning("ignoring: %s with %s architecture", so_name, arch)
continue
except ValueError:
shared_libraries_with_invalid_machine.append(so_name)
log.warning("ignoring: %s with unknown architecture", so_name)
continue

platform_wheel = True

for key, value in elf_find_versioned_symbols(elf):
log.debug("key %s, value %s", key, value)
versioned_symbols[key].add(value)

is_py_ext, py_ver = elf_is_python_extension(fn, elf)

# If the ELF is a Python extention, we definitely need to
# include its external dependencies.
if is_py_ext:
full_elftree[fn] = elftree
uses_PyFPE_jbuf |= elf_references_PyFPE_jbuf(elf)
if py_ver == 2:
uses_ucs2_symbols |= any(
True for _ in elf_find_ucs2_symbols(elf)
)
full_external_refs[fn] = wheel_policy.lddtree_external_references(
elftree, ctx.path
)
else:
# If the ELF is not a Python extension, it might be
# included in the wheel already because auditwheel repair
# vendored it, so we will check whether we should include
# its internal references later.
nonpy_elftree[fn] = elftree
POOL.submit(fn, inner, fn, elf)

# If at least one shared library exists in purelib, raise an error
if shared_libraries_in_purelib:
Expand All @@ -164,6 +176,8 @@ def get_wheel_elfdata(
wheel_policy.architecture, shared_libraries_with_invalid_machine
)

POOL.wait()

# Get a list of all external libraries needed by ELFs in the wheel.
needed_libs = {
lib
Expand Down

0 comments on commit 7d8abfd

Please sign in to comment.