Merge pull request #3756 from cmatsuoka/craft-1087/fix-parallel-build…

…-count parts: fix parallel build count
canonical · May 24, 2022 · e3d8149 · e3d8149
2 parents d64d2fb + 354ad89
commit e3d8149
Show file tree

Hide file tree

Showing 6 changed files with 109 additions and 5 deletions.
diff --git a/snapcraft/parts/lifecycle.py b/snapcraft/parts/lifecycle.py
@@ -142,14 +142,16 @@ def run(command_name: str, parsed_args: "argparse.Namespace") -> None:
     callbacks.register_prologue(_set_global_environment)
     callbacks.register_pre_step(_set_step_environment)
 
-    _expand_environment(yaml_data)
+    build_count = utils.get_parallel_build_count()
+    _expand_environment(yaml_data, parallel_build_count=build_count)
 
     project = Project.unmarshal(yaml_data)
 
     _run_command(
         command_name,
         project=project,
         parse_info=parse_info,
+        parallel_build_count=build_count,
         assets_dir=snap_project.assets_dir,
         parsed_args=parsed_args,
     )
@@ -161,6 +163,7 @@ def _run_command(
     project: Project,
     parse_info: Dict[str, List[str]],
     assets_dir: Path,
+    parallel_build_count: int,
     parsed_args: "argparse.Namespace",
 ) -> None:
     managed_mode = utils.is_managed_mode()
@@ -199,6 +202,7 @@ def _run_command(
         work_dir=work_dir,
         assets_dir=assets_dir,
         package_repositories=project.package_repositories,
+        parallel_build_count=parallel_build_count,
         part_names=part_names,
         adopt_info=project.adopt_info,
         project_name=project.name,
@@ -374,7 +378,7 @@ def _set_step_environment(step_info: StepInfo) -> bool:
     return True
 
 
-def _expand_environment(snapcraft_yaml: Dict[str, Any], *, parallel_build_count: int = 1) -> None:
+def _expand_environment(snapcraft_yaml: Dict[str, Any], *, parallel_build_count: int) -> None:
     """Expand global variables in the provided dictionary values.
 
     :param snapcraft_yaml: A dictionary containing the contents of the

diff --git a/snapcraft/parts/parts.py b/snapcraft/parts/parts.py
@@ -56,6 +56,7 @@ def __init__(
         work_dir: pathlib.Path,
         assets_dir: pathlib.Path,
         package_repositories: List[Dict[str, Any]],
+        parallel_build_count: int,
         part_names: Optional[List[str]],
         adopt_info: Optional[str],
         parse_info: Dict[str, List[str]],
@@ -89,6 +90,7 @@ def __init__(
                 ignore_local_sources=["*.snap"],
                 extra_build_packages=extra_build_packages,
                 extra_build_snaps=extra_build_snaps,
+                parallel_build_count=parallel_build_count,
                 project_name=project_name,
                 project_vars_part_name=adopt_info,
                 project_vars=project_vars,

diff --git a/snapcraft/utils.py b/snapcraft/utils.py
@@ -16,6 +16,7 @@
 
 """Utilities for snapcraft."""
 
+import multiprocessing
 import os
 import pathlib
 import platform
@@ -167,6 +168,41 @@ def get_effective_base(
     return name if project_type == "base" else base
 
 
+def get_parallel_build_count() -> int:
+    """Obtain the number of concurrent jobs to execute.
+
+    Try different strategies to obtain the number of parallel jobs
+    to execute. If they fail, assume the safe default of 1. The
+    number of concurrent build jobs can be limited by setting the
+    environment variable ``SNAPCRAFT_MAX_PARALLEL_BUILD_COUNT``.
+
+    :return: The number of parallel jobs for the current host.
+    """
+    try:
+        build_count = len(os.sched_getaffinity(0))
+        emit.trace(f"CPU count (from process affinity): {build_count}")
+    except AttributeError:
+        # Fall back to multiprocessing.cpu_count()...
+        try:
+            build_count = multiprocessing.cpu_count()
+            emit.trace(f"CPU count (from multiprocessing): {build_count}")
+        except NotImplementedError:
+            emit.message(
+                "Unable to determine CPU count; disabling parallel builds",
+                intermediate=True,
+            )
+            build_count = 1
+
+    try:
+        max_count = int(os.environ.get("SNAPCRAFT_MAX_PARALLEL_BUILD_COUNT", ""))
+        if max_count > 0:
+            build_count = min(build_count, max_count)
+    except ValueError:
+        emit.trace("Invalid SNAPCRAFT_MAX_PARALLEL_BUILD_COUNT value")
+
+    return build_count
+
+
 def confirm_with_user(prompt_text, default=False) -> bool:
     """Query user for yes/no answer.
 

diff --git a/tests/unit/parts/test_lifecycle.py b/tests/unit/parts/test_lifecycle.py
@@ -131,6 +131,7 @@ def test_snapcraft_yaml_load(new_dir, snapcraft_yaml, filename, mocker):
     """Snapcraft.yaml should be parsed as a valid yaml file."""
     yaml_data = snapcraft_yaml(base="core22", filename=filename)
     run_command_mock = mocker.patch("snapcraft.parts.lifecycle._run_command")
+    mocker.patch("snapcraft.utils.get_parallel_build_count", return_value=5)
 
     parts_lifecycle.run(
         "pull",
@@ -152,6 +153,7 @@ def test_snapcraft_yaml_load(new_dir, snapcraft_yaml, filename, mocker):
             project=project,
             parse_info={},
             assets_dir=assets_dir,
+            parallel_build_count=5,
             parsed_args=argparse.Namespace(
                 parts=["part1"], destructive_mode=True, use_lxd=False, provider=None
             ),
@@ -232,7 +234,12 @@ def test_lifecycle_run_command_step(
         setattr(parsed_args, debug_shell, True)
 
     parts_lifecycle._run_command(
-        cmd, project=project, parse_info={}, assets_dir=Path(), parsed_args=parsed_args
+        cmd,
+        project=project,
+        parse_info={},
+        assets_dir=Path(),
+        parallel_build_count=8,
+        parsed_args=parsed_args,
     )
 
     call_args = {"debug": False, "shell": False, "shell_after": False}
@@ -255,6 +262,7 @@ def test_lifecycle_run_command_pack(cmd, snapcraft_yaml, project_vars, new_dir,
         project=project,
         parse_info={},
         assets_dir=Path(),
+        parallel_build_count=8,
         parsed_args=argparse.Namespace(
             directory=None,
             output=None,
@@ -295,6 +303,7 @@ def test_lifecycle_pack_destructive_mode(
         project=project,
         parse_info={},
         assets_dir=Path(),
+        parallel_build_count=8,
         parsed_args=argparse.Namespace(
             directory=None,
             output=None,
@@ -334,6 +343,7 @@ def test_lifecycle_pack_managed(cmd, snapcraft_yaml, project_vars, new_dir, mock
         project=project,
         parse_info={},
         assets_dir=Path(),
+        parallel_build_count=8,
         parsed_args=argparse.Namespace(
             directory=None,
             output=None,
@@ -367,6 +377,7 @@ def test_lifecycle_pack_not_managed(cmd, snapcraft_yaml, new_dir, mocker):
         project=project,
         parse_info={},
         assets_dir=Path(),
+        parallel_build_count=8,
         parsed_args=argparse.Namespace(
             directory=None,
             output=None,
@@ -415,6 +426,7 @@ def test_lifecycle_pack_metadata_error(cmd, snapcraft_yaml, new_dir, mocker):
             project=project,
             assets_dir=Path(),
             parse_info={},
+            parallel_build_count=8,
             parsed_args=argparse.Namespace(
                 directory=None,
                 output=None,
@@ -469,6 +481,7 @@ def test_lifecycle_run_command_clean(snapcraft_yaml, project_vars, new_dir, mock
         project=project,
         parse_info={},
         assets_dir=Path(),
+        parallel_build_count=8,
         parsed_args=argparse.Namespace(
             directory=None,
             output=None,
@@ -493,6 +506,7 @@ def test_lifecycle_clean_destructive_mode(
         project=project,
         parse_info={},
         assets_dir=Path(),
+        parallel_build_count=8,
         parsed_args=argparse.Namespace(
             directory=None,
             output=None,
@@ -515,6 +529,7 @@ def test_lifecycle_clean_part_names(snapcraft_yaml, project_vars, new_dir, mocke
         project=project,
         parse_info={},
         assets_dir=Path(),
+        parallel_build_count=8,
         parsed_args=argparse.Namespace(
             directory=None,
             output=None,
@@ -551,6 +566,7 @@ def test_lifecycle_clean_part_names_destructive_mode(
         project=project,
         parse_info={},
         assets_dir=Path(),
+        parallel_build_count=8,
         parsed_args=argparse.Namespace(
             directory=None,
             output=None,
@@ -578,6 +594,7 @@ def test_lifecycle_clean_managed(snapcraft_yaml, project_vars, new_dir, mocker):
         project=project,
         parse_info={},
         assets_dir=Path(),
+        parallel_build_count=8,
         parsed_args=argparse.Namespace(
             directory=None,
             output=None,
@@ -604,6 +621,7 @@ def test_lifecycle_debug_shell(snapcraft_yaml, cmd, new_dir, mocker):
             project=project,
             parse_info={},
             assets_dir=Path(),
+            parallel_build_count=8,
             parsed_args=argparse.Namespace(
                 directory=None,
                 output=None,
@@ -637,6 +655,7 @@ def _fake_execute(_, action: Action, **kwargs):  # pylint: disable=unused-argume
         project=project,
         parse_info={},
         assets_dir=Path(),
+        parallel_build_count=8,
         parsed_args=argparse.Namespace(
             directory=None,
             output=None,
@@ -679,6 +698,7 @@ def _fake_execute(_, action: Action, **kwargs):  # pylint: disable=unused-argume
         project=project,
         parse_info={},
         assets_dir=Path(),
+        parallel_build_count=8,
         parsed_args=argparse.Namespace(
             directory=None,
             output=None,
@@ -806,7 +826,7 @@ def test_expand_environment(new_dir, mocker):
         },
         "field12": ["$CRAFT_PARALLEL_BUILD_COUNT", "$SNAPCRAFT_PARALLEL_BUILD_COUNT"],
     }
-    parts_lifecycle._expand_environment(yaml_data)
+    parts_lifecycle._expand_environment(yaml_data, parallel_build_count=8)
 
     assert yaml_data == {
         "name": "test-env",
@@ -825,7 +845,7 @@ def test_expand_environment(new_dir, mocker):
             "field10": [f"{new_dir}/prime", f"{new_dir}/prime"],
             "field11": [f"{new_dir}", f"{new_dir}"],
         },
-        "field12": ["1", "1"],
+        "field12": ["8", "8"],
     }
 
 

diff --git a/tests/unit/parts/test_parts.py b/tests/unit/parts/test_parts.py
@@ -39,6 +39,7 @@ def test_parts_lifecycle_run(mocker, parts_data, step_name, new_dir, emitter):
         parts_data,
         work_dir=new_dir,
         assets_dir=new_dir,
+        parallel_build_count=8,
         part_names=[],
         package_repositories=[],
         adopt_info=None,
@@ -59,6 +60,7 @@ def test_parts_lifecycle_run(mocker, parts_data, step_name, new_dir, emitter):
             ignore_local_sources=["*.snap"],
             extra_build_packages=[],
             extra_build_snaps=["core22"],
+            parallel_build_count=8,
             project_name="test-project",
             project_vars_part_name=None,
             project_vars={"version": "1", "grade": "stable"},
@@ -72,6 +74,7 @@ def test_parts_lifecycle_run_bad_step(parts_data, new_dir):
         parts_data,
         work_dir=new_dir,
         assets_dir=new_dir,
+        parallel_build_count=8,
         part_names=[],
         package_repositories=[],
         adopt_info=None,
@@ -89,6 +92,7 @@ def test_parts_lifecycle_run_internal_error(parts_data, new_dir, mocker):
         parts_data,
         work_dir=new_dir,
         assets_dir=new_dir,
+        parallel_build_count=8,
         part_names=[],
         package_repositories=[],
         adopt_info=None,
@@ -107,6 +111,7 @@ def test_parts_lifecycle_run_parts_error(new_dir):
         {"p1": {"plugin": "dump", "source": "foo"}},
         work_dir=new_dir,
         assets_dir=new_dir,
+        parallel_build_count=8,
         part_names=[],
         package_repositories=[],
         adopt_info=None,
@@ -126,6 +131,7 @@ def test_parts_lifecycle_clean(parts_data, new_dir, emitter):
         parts_data,
         work_dir=new_dir,
         assets_dir=new_dir,
+        parallel_build_count=8,
         part_names=[],
         package_repositories=[],
         adopt_info=None,
@@ -142,6 +148,7 @@ def test_parts_lifecycle_clean_parts(parts_data, new_dir, emitter):
         parts_data,
         work_dir=new_dir,
         assets_dir=new_dir,
+        parallel_build_count=8,
         part_names=[],
         package_repositories=[],
         adopt_info=None,
@@ -163,6 +170,7 @@ def test_parts_lifecycle_initialize_with_package_repositories(
         parts_data,
         work_dir=new_dir,
         assets_dir=new_dir,
+        parallel_build_count=8,
         part_names=[],
         package_repositories=[
             {
@@ -185,6 +193,7 @@ def test_parts_lifecycle_initialize_with_package_repositories(
             ignore_local_sources=["*.snap"],
             extra_build_packages=["gnupg", "dirmngr"],
             extra_build_snaps=["core22"],
+            parallel_build_count=8,
             project_name="test-project",
             project_vars_part_name=None,
             project_vars={"version": "1", "grade": "stable"},

diff --git a/tests/unit/test_utils.py b/tests/unit/test_utils.py
@@ -14,6 +14,7 @@
 # You should have received a copy of the GNU General Public License
 # along with this program.  If not, see <http://www.gnu.org/licenses/>.
 
+import os
 from textwrap import dedent
 
 import pytest
@@ -220,6 +221,38 @@ def test_get_host_architecture(
     assert utils.get_host_architecture() == deb_arch
 
 
+########################
+# Parallel build count #
+########################
+
+
+def test_get_parallel_build_count(mocker):
+    mocker.patch("os.sched_getaffinity", return_value=[1] * 13)
+    assert utils.get_parallel_build_count() == 13
+
+
+def test_get_parallel_build_count_no_affinity(mocker):
+    mocker.patch("os.sched_getaffinity", side_effect=AttributeError)
+    mocker.patch("multiprocessing.cpu_count", return_value=17)
+    assert utils.get_parallel_build_count() == 17
+
+
+def test_get_parallel_build_count_disable(mocker):
+    mocker.patch("os.sched_getaffinity", side_effect=AttributeError)
+    mocker.patch("multiprocessing.cpu_count", side_effect=NotImplementedError)
+    assert utils.get_parallel_build_count() == 1
+
+
+@pytest.mark.parametrize(
+    "max_count,count",
+    [("", 13), ("xx", 13), ("0", 13), ("1", 1), ("8", 8), ("13", 13), ("14", 13)],
+)
+def test_get_parallel_build_count_limited(mocker, max_count, count):
+    mocker.patch("os.sched_getaffinity", return_value=[1] * 13)
+    mocker.patch.dict(os.environ, {"SNAPCRAFT_MAX_PARALLEL_BUILD_COUNT": max_count})
+    assert utils.get_parallel_build_count() == count
+
+
 #################
 # Humanize List #
 #################