Skip to content

Commit 9d7ca27

Browse files
add cli arg to limit download parallelism
1 parent 90703cc commit 9d7ca27

File tree

8 files changed

+42
-3
lines changed

8 files changed

+42
-3
lines changed

src/pip/_internal/cli/cmdoptions.py

+18
Original file line numberDiff line numberDiff line change
@@ -231,6 +231,24 @@ class PipOption(Option):
231231
help="Specify whether the progress bar should be used [on, off, raw] (default: on)",
232232
)
233233

234+
235+
batch_download_parallelism: Callable[..., Option] = partial(
236+
Option,
237+
"--batch-download-parallelism",
238+
dest="batch_download_parallelism",
239+
type="int",
240+
default=10,
241+
help=(
242+
"Maximum parallelism employed for batch downloading of metadata-only dists"
243+
" (default %default parallel requests)."
244+
" Note that more than 10 downloads may overflow the requests connection pool,"
245+
" which may affect performance."
246+
" Note also that commands such as 'install --dry-run' should avoid downloads"
247+
" entirely, and so will not be affected by this option."
248+
),
249+
)
250+
251+
234252
log: Callable[..., Option] = partial(
235253
PipOption,
236254
"--log",

src/pip/_internal/cli/req_command.py

+2
Original file line numberDiff line numberDiff line change
@@ -100,6 +100,7 @@ def make_requirement_preparer(
100100
use_user_site: bool,
101101
download_dir: Optional[str] = None,
102102
verbosity: int = 0,
103+
batch_download_parallelism: Optional[int] = None,
103104
) -> RequirementPreparer:
104105
"""
105106
Create a RequirementPreparer instance for the given parameters.
@@ -141,6 +142,7 @@ def make_requirement_preparer(
141142
use_user_site=use_user_site,
142143
lazy_wheel=lazy_wheel,
143144
verbosity=verbosity,
145+
batch_download_parallelism=batch_download_parallelism,
144146
legacy_resolver=legacy_resolver,
145147
)
146148

src/pip/_internal/commands/download.py

+2
Original file line numberDiff line numberDiff line change
@@ -47,6 +47,7 @@ def add_options(self) -> None:
4747
self.cmd_opts.add_option(cmdoptions.pre())
4848
self.cmd_opts.add_option(cmdoptions.require_hashes())
4949
self.cmd_opts.add_option(cmdoptions.progress_bar())
50+
self.cmd_opts.add_option(cmdoptions.batch_download_parallelism())
5051
self.cmd_opts.add_option(cmdoptions.no_build_isolation())
5152
self.cmd_opts.add_option(cmdoptions.use_pep517())
5253
self.cmd_opts.add_option(cmdoptions.no_use_pep517())
@@ -116,6 +117,7 @@ def run(self, options: Values, args: List[str]) -> int:
116117
download_dir=options.download_dir,
117118
use_user_site=False,
118119
verbosity=self.verbosity,
120+
batch_download_parallelism=options.batch_download_parallelism,
119121
)
120122

121123
resolver = self.make_resolver(

src/pip/_internal/commands/install.py

+2
Original file line numberDiff line numberDiff line change
@@ -237,6 +237,7 @@ def add_options(self) -> None:
237237
self.cmd_opts.add_option(cmdoptions.prefer_binary())
238238
self.cmd_opts.add_option(cmdoptions.require_hashes())
239239
self.cmd_opts.add_option(cmdoptions.progress_bar())
240+
self.cmd_opts.add_option(cmdoptions.batch_download_parallelism())
240241
self.cmd_opts.add_option(cmdoptions.root_user_action())
241242

242243
index_opts = cmdoptions.make_option_group(
@@ -359,6 +360,7 @@ def run(self, options: Values, args: List[str]) -> int:
359360
finder=finder,
360361
use_user_site=options.use_user_site,
361362
verbosity=self.verbosity,
363+
batch_download_parallelism=options.batch_download_parallelism,
362364
)
363365
resolver = self.make_resolver(
364366
preparer=preparer,

src/pip/_internal/commands/wheel.py

+2
Original file line numberDiff line numberDiff line change
@@ -67,6 +67,7 @@ def add_options(self) -> None:
6767
self.cmd_opts.add_option(cmdoptions.ignore_requires_python())
6868
self.cmd_opts.add_option(cmdoptions.no_deps())
6969
self.cmd_opts.add_option(cmdoptions.progress_bar())
70+
self.cmd_opts.add_option(cmdoptions.batch_download_parallelism())
7071

7172
self.cmd_opts.add_option(
7273
"--no-verify",
@@ -131,6 +132,7 @@ def run(self, options: Values, args: List[str]) -> int:
131132
download_dir=options.wheel_dir,
132133
use_user_site=False,
133134
verbosity=self.verbosity,
135+
batch_download_parallelism=options.batch_download_parallelism,
134136
)
135137

136138
resolver = self.make_resolver(

src/pip/_internal/network/download.py

+11-2
Original file line numberDiff line numberDiff line change
@@ -13,7 +13,7 @@
1313
from pip._vendor.requests.models import Response
1414

1515
from pip._internal.cli.progress_bars import get_download_progress_renderer
16-
from pip._internal.exceptions import NetworkConnectionError
16+
from pip._internal.exceptions import CommandError, NetworkConnectionError
1717
from pip._internal.models.index import PyPI
1818
from pip._internal.models.link import Link
1919
from pip._internal.network.cache import is_from_cache
@@ -226,11 +226,20 @@ def __init__(
226226
self,
227227
session: PipSession,
228228
progress_bar: str,
229+
max_parallelism: Optional[int] = None,
229230
) -> None:
230231
self._session = session
231232
# FIXME: support progress bar with parallel downloads!
232233
logger.info("Ignoring progress bar %s for parallel downloads", progress_bar)
233234

235+
if max_parallelism is None:
236+
max_parallelism = 1
237+
if max_parallelism < 1:
238+
raise CommandError(
239+
f"invalid batch download parallelism {max_parallelism}: must be >=1"
240+
)
241+
self._max_parallelism: int = max_parallelism
242+
234243
def __call__(
235244
self, links: Iterable[Link], location: Path
236245
) -> Iterable[Tuple[Link, Tuple[Path, Optional[str]]]]:
@@ -254,7 +263,7 @@ def __call__(
254263
q: "Queue[Union[Tuple[Link, Path, Optional[str]], BaseException]]" = Queue()
255264
event = Event()
256265
# Limit downloads to 10 at a time so we can reuse our connection pool.
257-
semaphore = Semaphore(value=10)
266+
semaphore = Semaphore(value=self._max_parallelism)
258267

259268
# Distribute request i/o across equivalent threads.
260269
# NB: event-based/async is likely a better model than thread-per-request, but

src/pip/_internal/operations/prepare.py

+4-1
Original file line numberDiff line numberDiff line change
@@ -230,6 +230,7 @@ def __init__(
230230
use_user_site: bool,
231231
lazy_wheel: bool,
232232
verbosity: int,
233+
batch_download_parallelism: Optional[int],
233234
legacy_resolver: bool,
234235
) -> None:
235236
super().__init__()
@@ -239,7 +240,9 @@ def __init__(
239240
self.build_tracker = build_tracker
240241
self._session = session
241242
self._download = Downloader(session, progress_bar)
242-
self._batch_download = BatchDownloader(session, progress_bar)
243+
self._batch_download = BatchDownloader(
244+
session, progress_bar, max_parallelism=batch_download_parallelism
245+
)
243246
self.finder = finder
244247

245248
# Where still-packed archives should be written to. If None, they are

tests/unit/test_req.py

+1
Original file line numberDiff line numberDiff line change
@@ -106,6 +106,7 @@ def _basic_resolver(
106106
use_user_site=False,
107107
lazy_wheel=False,
108108
verbosity=0,
109+
batch_download_parallelism=None,
109110
legacy_resolver=True,
110111
)
111112
yield Resolver(

0 commit comments

Comments
 (0)