Merge pull request #392 from tiran/build-settings-jobs

Redesign parallel build configuration

Author: mergify[bot]

requirements-mypy.txt

@@ -1,6 +1,7 @@
   mypy
   pytest
   types-html5lib
+  types-psutil
   types-PyYAML
   types-requests
   types-toml

requirements.txt

@@ -3,6 +3,7 @@ elfdeps>=0.2.0
 html5lib
 packaging
 pkginfo
+psutil
 pydantic
 PyGithub
 pyproject_hooks>=1.0.0,!=1.1.0

src/fromager/__main__.py

@@ -1,7 +1,6 @@
 #!/usr/bin/env python3
 
 import logging
-import os
 import pathlib
 
 import click
@@ -119,7 +118,8 @@
     "-j",
     "--jobs",
     type=int,
-    help="number of jobs available to run in parallel",
+    default=None,
+    help="maximum number of jobs to run in parallel",
 )
 @click.option(
     "--network-isolation/--no-network-isolation",
@@ -143,7 +143,7 @@ def main(
     wheel_server_url: str,
     cleanup: bool,
     variant: str,
-    jobs: int,
+    jobs: int | None,
     network_isolation: bool,
 ) -> None:
     # Set the overall logger level to debug and allow the handlers to filter
@@ -190,6 +190,7 @@ def main(
             settings_dir=settings_dir,
             patches_dir=patches_dir,
             variant=variant,
+            max_jobs=jobs,
         ),
         constraints_file=constraints_file,
         patches_dir=patches_dir,
@@ -199,8 +200,8 @@ def main(
         wheel_server_url=wheel_server_url,
         cleanup=cleanup,
         variant=variant,
-        jobs=jobs if jobs is None or jobs > 0 else os.cpu_count(),
         network_isolation=network_isolation,
+        max_jobs=jobs,
     )
     wkctx.setup()
     ctx.obj = wkctx

src/fromager/context.py

@@ -36,15 +36,16 @@ def __init__(
         wheel_server_url: str,
         cleanup: bool = True,
         variant: str = "cpu",
-        jobs: int | None = None,
         network_isolation: bool = False,
+        max_jobs: int | None = None,
     ):
         if active_settings is None:
             active_settings = packagesettings.Settings(
                 settings=packagesettings.SettingsFile(),
                 package_settings=[],
                 patches_dir=patches_dir,
                 variant=variant,
+                max_jobs=max_jobs,
             )
         self.settings = active_settings
         self.input_constraints_file = constraints_file
@@ -64,7 +65,6 @@ def __init__(
         self.wheel_server_url = wheel_server_url
         self.cleanup = cleanup
         self.variant = variant
-        self.jobs = jobs
         self.network_isolation = network_isolation
 
         self._build_order_filename = self.work_dir / "build-order.json"

src/fromager/packagesettings.py

@@ -6,6 +6,7 @@
 import typing
 from collections.abc import Mapping
 
+import psutil
 import pydantic
 import yaml
 from packaging.utils import BuildTag, NormalizedName, canonicalize_name
@@ -142,6 +143,27 @@ def validate_destination_filename(cls, v):
         return v
 
 
+class BuildOptions(pydantic.BaseModel):
+    """Build system options"""
+
+    model_config = MODEL_CONFIG
+
+    cpu_cores_per_job: int = Field(default=1, ge=1)
+    """Scale parallel jobs by available CPU cores
+
+    Examples:
+    1: as many parallel jobs as CPU logical cores
+    2: allocate 2 cores per job
+    """
+
+    memory_per_job_gb: float = Field(default=1.0, ge=0.1)
+    """Scale parallel jobs by available virtual memory (without swap)
+
+    Examples:
+    0.5: assume each parallel job requires 512 MB virtual memory
+    """
+
+
 class VariantInfo(pydantic.BaseModel):
     """Variant information for a package"""
 
@@ -209,6 +231,9 @@ class PackageSettings(pydantic.BaseModel):
     resolve_source: ResolveSource = Field(default_factory=ResolveSource)
     """Resolve distribution version"""
 
+    build_options: BuildOptions = Field(default_factory=BuildOptions)
+    """Build system options"""
+
     variants: Mapping[Variant, VariantInfo] = Field(default_factory=dict)
     """Variant configuration"""
 
@@ -320,16 +345,30 @@ def _resolve_template(
         raise
 
 
+def get_cpu_count() -> int:
+    """CPU count from scheduler affinity"""
+    if hasattr(os, "sched_getaffinity"):
+        return len(os.sched_getaffinity(0))
+    else:
+        return os.cpu_count() or 1
+
+
+def get_available_memory_gib() -> float:
+    """available virtual memory in GiB"""
+    return psutil.virtual_memory().available / (1024**3)
+
+
 class PackageBuildInfo:
     """Package build information
 
     Public API for PackageSettings with i
     """
 
-    def __init__(self, ctx: "Settings", ps: PackageSettings) -> None:
-        self._variant = typing.cast(Variant, ctx.variant)
-        self._patches_dir = ctx.patches_dir
-        self._variant_changelog = ctx.variant_changelog()
+    def __init__(self, settings: "Settings", ps: PackageSettings) -> None:
+        self._variant = typing.cast(Variant, settings.variant)
+        self._patches_dir = settings.patches_dir
+        self._variant_changelog = settings.variant_changelog()
+        self._max_jobs: int | None = settings.max_jobs
         self._ps = ps
         self._plugin_module: types.ModuleType | None | typing.Literal[False] = False
         self._patches: PatchMap | None = None
@@ -497,6 +536,35 @@ def get_extra_environ(
 
         return extra_environ
 
+    def parallel_jobs(self) -> int:
+        """How many parallel jobs?"""
+        # adjust by CPU cores, at least 1
+        cpu_cores_per_job = self._ps.build_options.cpu_cores_per_job
+        cpu_count = get_cpu_count()
+        max_num_job_cores = int(max(1, cpu_count // cpu_cores_per_job))
+        logger.debug(
+            f"{self.package}: {max_num_job_cores=}, {cpu_cores_per_job=}, {cpu_count=}"
+        )
+
+        # adjust by memory consumption per job, at least 1
+        memory_per_job_gb = self._ps.build_options.memory_per_job_gb
+        free_memory = get_available_memory_gib()
+        max_num_jobs_memory = int(max(1.0, free_memory // memory_per_job_gb))
+        logger.debug(
+            f"{self.package}: {max_num_jobs_memory=}, {memory_per_job_gb=}, {free_memory=:0.1f} GiB"
+        )
+
+        # limit by smallest amount of CPU, memory, and --jobs parameter
+        max_jobs = cpu_count if self._max_jobs is None else self._max_jobs
+        parallel_builds = min(max_num_job_cores, max_num_jobs_memory, max_jobs)
+
+        logger.debug(
+            f"{self.package}: parallel builds {parallel_builds=} "
+            f"({free_memory=:0.1f} GiB, {cpu_count=}, {max_jobs=})"
+        )
+
+        return parallel_builds
+
     def serialize(self, **kwargs) -> dict[str, typing.Any]:
         return self._ps.serialize(**kwargs)
 
@@ -556,13 +624,15 @@ def __init__(
         package_settings: typing.Iterable[PackageSettings],
         variant: Variant | str,
         patches_dir: pathlib.Path,
+        max_jobs: int | None,
     ) -> None:
         self._settings = settings
         self._package_settings: dict[Package, PackageSettings] = {
             p.name: p for p in package_settings
         }
         self._variant = typing.cast(Variant, variant)
         self._patches_dir = patches_dir
+        self._max_jobs = max_jobs
         self._pbi_cache: dict[Package, PackageBuildInfo] = {}
 
     @classmethod
@@ -573,6 +643,7 @@ def from_files(
         settings_dir: pathlib.Path,
         variant: Variant | str,
         patches_dir: pathlib.Path,
+        max_jobs: int | None,
     ) -> "Settings":
         """Create Settings from settings.yaml and directory"""
         if settings_file.is_file():
@@ -591,6 +662,7 @@ def from_files(
             package_settings=package_settings,
             variant=variant,
             patches_dir=patches_dir,
+            max_jobs=max_jobs,
         )
 
     @property
@@ -612,10 +684,21 @@ def patches_dir(self) -> pathlib.Path:
 
     @patches_dir.setter
     def patches_dir(self, path: pathlib.Path) -> None:
-        """Change patches_dr (for testing)"""
+        """Change patches_dir (for testing)"""
         self._pbi_cache.clear()
         self._patches_dir = path
 
+    @property
+    def max_jobs(self) -> int | None:
+        """Get max parallel jobs"""
+        return self._max_jobs
+
+    @max_jobs.setter
+    def max_jobs(self, jobs: int | None) -> None:
+        """Change max jobs (for testing)"""
+        self._pbi_cache.clear()
+        self._max_jobs = jobs
+
     def variant_changelog(self) -> list[str]:
         """Get global changelog for current variant"""
         return list(self._settings.changelog.get(self.variant, []))

src/fromager/wheels.py

@@ -234,13 +234,13 @@ def build_wheel(
     # TODO: refactor?
     # Build Rust without network access
     extra_environ["CARGO_NET_OFFLINE"] = "true"
-    # configure max jobs settings. should cover most of the cases, if not then the user can use ctx.jobs in their plugin
-    if ctx.jobs:
-        extra_environ["MAKEFLAGS"] = (
-            f"{extra_environ.get('MAKEFLAGS', '')} -j{ctx.jobs}"
-        )
-        extra_environ["CMAKE_BUILD_PARALLEL_LEVEL"] = f"{ctx.jobs}"
-        extra_environ["MAX_JOBS"] = f"{ctx.jobs}"
+
+    # configure max jobs settings, settings depend on package, available
+    # CPU cores, and available virtual memory.
+    jobs = pbi.parallel_jobs()
+    extra_environ["MAKEFLAGS"] = f"{extra_environ.get('MAKEFLAGS', '')} -j{jobs}"
+    extra_environ["CMAKE_BUILD_PARALLEL_LEVEL"] = str(jobs)
+    extra_environ["MAX_JOBS"] = str(jobs)
 
     # Start the timer
     start = datetime.now().replace(microsecond=0)

tests/conftest.py

@@ -45,6 +45,7 @@ def testdata_context(
             settings_dir=overrides / "settings",
             patches_dir=patches_dir,
             variant=variant,
+            max_jobs=None,
         ),
         constraints_file=None,
         patches_dir=overrides / "patches",

tests/test_packagesettings.py

@@ -1,4 +1,5 @@
 import pathlib
+from unittest.mock import Mock, patch
 
 import pydantic
 import pytest
@@ -21,6 +22,10 @@
 
 FULL_EXPECTED = {
     "build_dir": pathlib.Path("python"),
+    "build_options": {
+        "cpu_cores_per_job": 4,
+        "memory_per_job_gb": 4.0,
+    },
     "changelog": {
         Version("1.0.1"): ["fixed bug"],
         Version("1.0.2"): ["more bugs", "rebuild"],
@@ -64,6 +69,10 @@
 EMPTY_EXPECTED = {
     "name": "test-empty-pkg",
     "build_dir": None,
+    "build_options": {
+        "cpu_cores_per_job": 1,
+        "memory_per_job_gb": 1.0,
+    },
     "changelog": {},
     "env": {},
     "download_source": {
@@ -302,3 +311,45 @@ def test_settings_list(testdata_context: context.WorkContext) -> None:
     assert testdata_context.settings.variant_changelog() == [
         "setuptools upgraded to 82.0.0"
     ]
+
+
+@patch("fromager.packagesettings.get_cpu_count", return_value=8)
+@patch("fromager.packagesettings.get_available_memory_gib", return_value=7.1)
+def test_parallel_jobs(
+    get_available_memory_gib: Mock,
+    get_cpu_count: Mock,
+    testdata_context: context.WorkContext,
+) -> None:
+    assert testdata_context.settings.max_jobs is None
+
+    pbi = testdata_context.settings.package_build_info(TEST_EMPTY_PKG)
+    assert pbi.parallel_jobs() == 7
+
+    get_cpu_count.return_value = 4
+    assert pbi.parallel_jobs() == 4
+
+    get_available_memory_gib.return_value = 2.1
+    assert pbi.parallel_jobs() == 2
+
+    get_available_memory_gib.return_value = 1.5
+    assert pbi.parallel_jobs() == 1
+
+    testdata_context.settings.max_jobs = 2
+    pbi = testdata_context.settings.package_build_info(TEST_EMPTY_PKG)
+    get_available_memory_gib.return_value = 23
+    assert pbi.parallel_jobs() == 2
+
+    # test-pkg needs more memory
+    testdata_context.settings.max_jobs = 200
+    pbi = testdata_context.settings.package_build_info(TEST_PKG)
+    get_cpu_count.return_value = 16
+    get_available_memory_gib.return_value = 20
+    assert pbi.parallel_jobs() == 4
+
+    get_cpu_count.return_value = 32
+    get_available_memory_gib.return_value = 25
+    assert pbi.parallel_jobs() == 6
+
+    testdata_context.settings.max_jobs = 4
+    pbi = testdata_context.settings.package_build_info(TEST_PKG)
+    assert pbi.parallel_jobs() == 4

tests/testdata/context/overrides/settings/test_pkg.yaml

@@ -1,4 +1,7 @@
 build_dir: python
+build_options:
+    cpu_cores_per_job: 4
+    memory_per_job_gb: 4
 changelog:
     "1.0.1":
         - fixed bug