diff --git a/src/pip/_internal/req/__init__.py b/src/pip/_internal/req/__init__.py index 422d851d729..db3db5311fe 100644 --- a/src/pip/_internal/req/__init__.py +++ b/src/pip/_internal/req/__init__.py @@ -1,5 +1,8 @@ import collections import logging +import os +import threading +from concurrent.futures import ThreadPoolExecutor, as_completed from dataclasses import dataclass from typing import Generator, List, Optional, Sequence, Tuple @@ -32,6 +35,48 @@ def _validate_requirements( yield req.name, req +def install_requirement( + req_name, + requirement, + global_options, + root, + home, + prefix, + warn_script_location, + use_user_site, + pycompile, +): + if requirement.should_reinstall: + logger.info("Attempting uninstall: %s", req_name) + with indent_log(): + uninstalled_pathset = requirement.uninstall(auto_confirm=True) + else: + uninstalled_pathset = None + + try: + print(f"installing [{threading.get_ident()}]: {req_name}") + requirement.install( + global_options, + root=root, + home=home, + prefix=prefix, + warn_script_location=warn_script_location, + use_user_site=use_user_site, + pycompile=pycompile, + ) + print(f" done [{threading.get_ident()}]: {req_name}") + except Exception: + # if install did not succeed, rollback previous uninstall + if uninstalled_pathset and not requirement.install_succeeded: + uninstalled_pathset.rollback() + raise + else: + if uninstalled_pathset and requirement.install_succeeded: + uninstalled_pathset.commit() + + return req_name + + def install_given_reqs( requirements: List[InstallRequirement], global_options: Sequence[str], @@ -56,35 +101,44 @@ def install_given_reqs( ) installed = [] + import time + start = time.perf_counter() + exception = None with indent_log(): - for req_name, requirement in to_install.items(): - if requirement.should_reinstall: - logger.info("Attempting uninstall: %s", req_name) - with indent_log(): - uninstalled_pathset = requirement.uninstall(auto_confirm=True) - else: - uninstalled_pathset = None - - try: - requirement.install( + workers = int(os.environ.get("PIP_THREADS", 1)) + with ThreadPoolExecutor( + max_workers=workers, thread_name_prefix="install_given_reqs" + ) as executor: + futures = { + executor.submit( + install_requirement, + req_name, + requirement, global_options, - root=root, - home=home, - prefix=prefix, - warn_script_location=warn_script_location, - use_user_site=use_user_site, - pycompile=pycompile, - ) - except Exception: - # if install did not succeed, rollback previous uninstall - if uninstalled_pathset and not requirement.install_succeeded: - uninstalled_pathset.rollback() - raise - else: - if uninstalled_pathset and requirement.install_succeeded: - uninstalled_pathset.commit() - - installed.append(InstallationResult(req_name)) + root, + home, + prefix, + warn_script_location, + use_user_site, + pycompile, + ): req_name + for req_name, requirement in to_install.items() + } + for future in as_completed(futures): + req_name = futures[future] + try: + result = future.result() + installed.append(InstallationResult(result)) + except Exception as e: + logger.error("Installation failed for %s: %s", req_name, str(e)) + # COMPATIBILITY: cancel_futures was added in python 3.9 + executor.shutdown(wait=False, cancel_futures=True) + exception = e + end = time.perf_counter() + print(f"total extraction time: {end-start:.3f}") + + if exception is not None: + raise exception return installed