Add maximum of workers to config (#3516)

.gitignore

@@ -38,3 +38,5 @@ MANIFEST.in
 /releases/*
 pip-wheel-metadata
 /poetry.toml
+
+poetry/core/*

docs/configuration.md

@@ -123,6 +123,21 @@ Defaults to one of the following directories:
 Use parallel execution when using the new (`>=1.1.0`) installer.
 Defaults to `true`.
 
+### `installer.max-workers`
+
+**Type**: int
+
+Set the maximum number of workers while using the parallel installer. Defaults to `number_of_cores + 4`.
+The `number_of_cores` is determined by `os.cpu_count()`.
+If this raises a `NotImplentedError` exception `number_of_cores` is assumed to be 1.
+
+If this configuration parameter is set to a value greater than `number_of_cores + 4`,
+the number of maximum workers is still limited at `number_of_cores + 4`.
+
+{{% note %}}
+This configuration will be ignored when `installer.parallel` is set to false.
+{{% /note %}}
+
 ### `virtualenvs.create`
 
 **Type**: boolean

pyproject.toml

@@ -65,6 +65,7 @@ deepdiff = "^5.0"
 httpretty = "^1.0"
 typing-extensions = { version = "^4.0.0", python = "<3.8" }
 zipp = { version = "^3.4", python = "<3.8" }
+flatdict = "^4.0.1"
 
 [tool.poetry.scripts]
 poetry = "poetry.console.application:main"

src/poetry/config/config.py

@@ -25,6 +25,10 @@ def boolean_normalizer(val: str) -> bool:
     return val in ["true", "1"]
 
 
+def int_normalizer(val: str) -> int:
+    return int(val)
+
+
 class Config:
 
     default_config: Dict[str, Any] = {
@@ -36,7 +40,7 @@ class Config:
             "options": {"always-copy": False, "system-site-packages": False},
         },
         "experimental": {"new-installer": True},
-        "installer": {"parallel": True},
+        "installer": {"parallel": True, "max-workers": None},
     }
 
     def __init__(
@@ -129,7 +133,8 @@ def process(self, value: Any) -> Any:
 
         return re.sub(r"{(.+?)}", lambda m: self.get(m.group(1)), value)
 
-    def _get_normalizer(self, name: str) -> Callable:
+    @staticmethod
+    def _get_normalizer(name: str) -> Callable:
         if name in {
             "virtualenvs.create",
             "virtualenvs.in-project",
@@ -143,4 +148,7 @@ def _get_normalizer(self, name: str) -> Callable:
         if name == "virtualenvs.path":
             return lambda val: str(Path(val))
 
+        if name == "installer.max-workers":
+            return int_normalizer
+
         return lambda val: val

src/poetry/console/commands/config.py

@@ -52,6 +52,7 @@ def unique_config_values(self) -> Dict[str, Tuple[Any, Any, Any]]:
 
         from poetry.config.config import boolean_normalizer
         from poetry.config.config import boolean_validator
+        from poetry.config.config import int_normalizer
         from poetry.locations import CACHE_DIR
 
         unique_config_values = {
@@ -87,6 +88,11 @@ def unique_config_values(self) -> Dict[str, Tuple[Any, Any, Any]]:
                 boolean_normalizer,
                 True,
             ),
+            "installer.max-workers": (
+                lambda val: int(val) > 0,
+                int_normalizer,
+                None,
+            ),
         }
 
         return unique_config_values

src/poetry/installation/executor.py

@@ -12,6 +12,7 @@
 from typing import Any
 from typing import Dict
 from typing import List
+from typing import Optional
 from typing import Union
 
 from cleo.io.null_io import NullIO
@@ -66,14 +67,9 @@ def __init__(
             parallel = config.get("installer.parallel", True)
 
         if parallel:
-            # This should be directly handled by ThreadPoolExecutor
-            # however, on some systems the number of CPUs cannot be determined
-            # (it raises a NotImplementedError), so, in this case, we assume
-            # that the system only has one CPU.
-            try:
-                self._max_workers = os.cpu_count() + 4
-            except NotImplementedError:
-                self._max_workers = 5
+            self._max_workers = self._get_max_workers(
+                desired_max_workers=config.get("installer.max-workers")
+            )
         else:
             self._max_workers = 1
 
@@ -190,6 +186,21 @@ def execute(self, operations: List["OperationTypes"]) -> int:
 
         return 1 if self._shutdown else 0
 
+    @staticmethod
+    def _get_max_workers(desired_max_workers: Optional[int] = None):
+        # This should be directly handled by ThreadPoolExecutor
+        # however, on some systems the number of CPUs cannot be determined
+        # (it raises a NotImplementedError), so, in this case, we assume
+        # that the system only has one CPU.
+        try:
+            default_max_workers = os.cpu_count() + 4
+        except NotImplementedError:
+            default_max_workers = 5
+
+        if desired_max_workers is None:
+            return default_max_workers
+        return min(default_max_workers, desired_max_workers)
+
     def _write(self, operation: "OperationTypes", line: str) -> None:
         if not self.supports_fancy_output() or not self._should_write_operation(
             operation

tests/config/test_config.py

@@ -2,31 +2,29 @@
 import re
 
 from typing import TYPE_CHECKING
-from typing import Any
-from typing import Dict
+from typing import Callable
 from typing import Iterator
-from typing import Optional
 from typing import Tuple
 
 import pytest
 
+from flatdict import FlatDict
+
 from poetry.config.config import Config
+from poetry.config.config import boolean_normalizer
+from poetry.config.config import int_normalizer
 
 
 if TYPE_CHECKING:
     from pathlib import Path
 
 
-def get_boolean_options(config: Optional[Dict[str, Any]] = None) -> str:
-    if config is None:
-        config = Config.default_config
+def get_options_based_on_normalizer(normalizer: Callable) -> str:
+    flattened_config = FlatDict(Config.default_config, delimiter=".")
 
-    for k, v in config.items():
-        if isinstance(v, bool) or v is None:
+    for k in flattened_config:
+        if Config._get_normalizer(k) == normalizer:
             yield k
-        if isinstance(v, dict):
-            for suboption in get_boolean_options(v):
-                yield f"{k}.{suboption}"
 
 
 @pytest.mark.parametrize(
@@ -43,10 +41,14 @@ def test_config_get_processes_depended_on_values(
 
 
 def generate_environment_variable_tests() -> Iterator[Tuple[str, str, str, bool]]:
-    for env_value, value in [("true", True), ("false", False)]:
-        for name in get_boolean_options():
-            env_var = "POETRY_{}".format(re.sub("[.-]+", "_", name).upper())
-            yield name, env_var, env_value, value
+    for normalizer, values in [
+        (boolean_normalizer, [("true", True), ("false", False)]),
+        (int_normalizer, [("4", 4), ("2", 2)]),
+    ]:
+        for env_value, value in values:
+            for name in get_options_based_on_normalizer(normalizer=normalizer):
+                env_var = "POETRY_" + re.sub("[.-]+", "_", name).upper()
+                yield name, env_var, env_value, value
 
 
 @pytest.mark.parametrize(

tests/console/commands/test_config.py

@@ -46,6 +46,7 @@ def test_list_displays_default_value_if_not_set(
 
     expected = """cache-dir = {cache}
 experimental.new-installer = true
+installer.max-workers = null
 installer.parallel = true
 virtualenvs.create = true
 virtualenvs.in-project = null
@@ -70,6 +71,7 @@ def test_list_displays_set_get_setting(
 
     expected = """cache-dir = {cache}
 experimental.new-installer = true
+installer.max-workers = null
 installer.parallel = true
 virtualenvs.create = false
 virtualenvs.in-project = null
@@ -118,6 +120,7 @@ def test_list_displays_set_get_local_setting(
 
     expected = """cache-dir = {cache}
 experimental.new-installer = true
+installer.max-workers = null
 installer.parallel = true
 virtualenvs.create = false
 virtualenvs.in-project = null

tests/installation/test_executor.py

@@ -559,3 +559,35 @@ def test_executor_should_use_cached_link_and_hash(
         Link("https://example.com/demo-0.1.0-py2.py3-none-any.whl"),
     )
     assert archive == link_cached
+
+
+@pytest.mark.parametrize(
+    ("max_workers", "cpu_count", "side_effect", "expected_workers"),
+    [
+        (None, 3, None, 7),
+        (3, 4, None, 3),
+        (8, 3, None, 7),
+        (None, 8, NotImplementedError(), 5),
+        (2, 8, NotImplementedError(), 2),
+        (8, 8, NotImplementedError(), 5),
+    ],
+)
+def test_executor_should_be_initialized_with_correct_workers(
+    tmp_venv,
+    pool,
+    config,
+    io,
+    mocker,
+    max_workers,
+    cpu_count,
+    side_effect,
+    expected_workers,
+):
+    config = Config()
+    config.merge({"installer": {"max-workers": max_workers}})
+
+    mocker.patch("os.cpu_count", return_value=cpu_count, side_effect=side_effect)
+
+    executor = Executor(tmp_venv, pool, config, io)
+
+    assert executor._max_workers == expected_workers