Allow control of multiprocessing using env var

This provides the ability to adjust the arbitrary cut off, or turn off
multiprocessing altogether.

Author: seddonym

CHANGELOG.rst

@@ -2,6 +2,12 @@
 Changelog
 =========
 
+latest
+------
+
+* Provide more control of multiprocessing via ``GRIMP_MIN_MULTIPROCESSING_MODULES``
+  environment variable.
+
 3.8.1 (2025-04-23)
 ------------------
 

docs/usage.rst

@@ -82,7 +82,12 @@ Building the graph
     :param str, optional cache_dir: The directory to use for caching the graph. Defaults to ``.grimp_cache``. To disable caching,
         pass ``None``. See :doc:`caching`.
     :return: An import graph that you can use to analyse the package.
-    :rtype: ImportGraph
+    :rtype: ``ImportGraph``
+
+    This method uses multiple operating system processes to build the graph, if the number of modules to scan (not
+    including modules in the cache) is 50 or more. This threshold can be adjusted by setting the ``GRIMP_MIN_MULTIPROCESSING_MODULES``
+    environment variable to a different number. To disable multiprocessing altogether, set it to a large number (more than
+    the number of modules in the codebase being analyzed).
 
 .. _typing module documentation: https://docs.python.org/3/library/typing.html#typing.TYPE_CHECKING
 

src/grimp/application/usecases.py

@@ -15,14 +15,19 @@
 from ..application.ports.packagefinder import AbstractPackageFinder
 from ..domain.valueobjects import DirectImport, Module
 from .config import settings
+import os
 
 
 class NotSupplied:
     pass
 
 
+# Calling code can set this environment variable if it wants to tune when to switch to
+# multiprocessing, or set it to a large number to disable it altogether.
+MIN_NUMBER_OF_MODULES_TO_SCAN_USING_MULTIPROCESSING_ENV_NAME = "GRIMP_MIN_MULTIPROCESSING_MODULES"
 # This is an arbitrary number, but setting it too low slows down our functional tests considerably.
-MIN_NUMBER_OF_MODULES_TO_SCAN_USING_MULTIPROCESSING = 50
+# If you change this, update docs/usage.rst too!
+DEFAULT_MIN_NUMBER_OF_MODULES_TO_SCAN_USING_MULTIPROCESSING = 50
 
 
 def build_graph(
@@ -238,7 +243,13 @@ def _create_chunks(module_files: Collection[ModuleFile]) -> tuple[tuple[ModuleFi
 
 
 def _decide_number_of_processes(number_of_module_files: int) -> int:
-    if number_of_module_files < MIN_NUMBER_OF_MODULES_TO_SCAN_USING_MULTIPROCESSING:
+    min_number_of_modules = int(
+        os.environ.get(
+            MIN_NUMBER_OF_MODULES_TO_SCAN_USING_MULTIPROCESSING_ENV_NAME,
+            DEFAULT_MIN_NUMBER_OF_MODULES_TO_SCAN_USING_MULTIPROCESSING,
+        )
+    )
+    if number_of_module_files < min_number_of_modules:
         # Don't incur the overhead of multiple processes.
         return 1
     return min(joblib.cpu_count(), number_of_module_files)

tests/functional/test_build_and_use_graph.py

@@ -3,7 +3,7 @@
 import pytest
 from unittest.mock import patch
 from grimp.application import usecases
-
+import os
 """
 For ease of reference, these are the imports of all the files:
 
@@ -55,7 +55,9 @@ def test_modules():
     }
 
 
-@patch.object(usecases, "MIN_NUMBER_OF_MODULES_TO_SCAN_USING_MULTIPROCESSING", 0)
+@patch.object(os, "environ", {
+    usecases.MIN_NUMBER_OF_MODULES_TO_SCAN_USING_MULTIPROCESSING_ENV_NAME: "0"
+})
 def test_modules_multiprocessing():
     """
     This test runs relatively slowly, but it's important we cover the multiprocessing code.

tests/unit/application/test_usecases.py

@@ -1,3 +1,4 @@
+import os
 from typing import Dict, Optional, Set
 from unittest.mock import sentinel, patch
 
@@ -13,6 +14,8 @@
 from tests.adaptors.modulefinder import BaseFakeModuleFinder
 from tests.config import override_settings
 
+SOME_CPU_COUNT = 8
+
 
 class TestBuildGraph:
     @pytest.mark.parametrize("include_external_packages", (True, False))
@@ -135,17 +138,49 @@ def write(
             usecases.build_graph("mypackage", **kwargs)
 
     @patch.object(usecases, "_scan_chunks", return_value={})
-    @patch.object(joblib, "cpu_count", return_value=8)
+    @patch.object(joblib, "cpu_count", return_value=SOME_CPU_COUNT)
     @pytest.mark.parametrize(
-        "number_of_modules, expected_number_of_chunks",
+        "number_of_modules, fake_environ, expected_number_of_chunks",
         [
-            (49, 1),  # Below threshold - just use one.
-            (50, 8),  # At threshold - use number of CPUs.
-            (1000, 8),  # Above threshold - use number of CPUs.
+            (
+                usecases.DEFAULT_MIN_NUMBER_OF_MODULES_TO_SCAN_USING_MULTIPROCESSING - 1,
+                {},
+                1,
+            ),
+            (
+                usecases.DEFAULT_MIN_NUMBER_OF_MODULES_TO_SCAN_USING_MULTIPROCESSING,
+                {},
+                SOME_CPU_COUNT,
+            ),
+            (
+                usecases.DEFAULT_MIN_NUMBER_OF_MODULES_TO_SCAN_USING_MULTIPROCESSING + 1,
+                {},
+                SOME_CPU_COUNT,
+            ),
+            (
+                149,
+                {usecases.MIN_NUMBER_OF_MODULES_TO_SCAN_USING_MULTIPROCESSING_ENV_NAME: 150},
+                1,
+            ),
+            (
+                150,
+                {usecases.MIN_NUMBER_OF_MODULES_TO_SCAN_USING_MULTIPROCESSING_ENV_NAME: 150},
+                SOME_CPU_COUNT,
+            ),
+            (
+                151,
+                {usecases.MIN_NUMBER_OF_MODULES_TO_SCAN_USING_MULTIPROCESSING_ENV_NAME: 150},
+                SOME_CPU_COUNT,
+            ),
         ],
     )
     def test_scanning_multiprocessing_respects_min_number_of_modules(
-        self, mock_cpu_count, mock_scan_chunks, number_of_modules, expected_number_of_chunks
+        self,
+        mock_cpu_count,
+        mock_scan_chunks,
+        number_of_modules,
+        fake_environ,
+        expected_number_of_chunks,
     ):
         class FakePackageFinder(BaseFakePackageFinder):
             directory_map = {"mypackage": "/path/to/mypackage"}
@@ -167,7 +202,7 @@ class FakeModuleFinder(BaseFakeModuleFinder):
             FILE_SYSTEM=FakeFileSystem(),
             PACKAGE_FINDER=FakePackageFinder(),
             MODULE_FINDER=FakeModuleFinder(),
-        ):
+        ), patch.object(os, "environ", fake_environ):
             usecases.build_graph("mypackage", cache_dir=None)
 
         [call] = mock_scan_chunks.call_args_list