fix(whl_library): avoid unnecessary repository rule restarts (bazel-contrib#1400)

aignas · web-flow · commit 6461a693a919 · 2023-09-04T05:25:16.000Z
Put the `PYTHONPATH` entries used in wheel building as a default value to a private attribute of the `whl_library` repository rule and use resolved path of the interpreter target in creating execution environment to avoid repository rule restarts when fetching external dependencies. The extra private attribute on the `whl_library` removes all but one restart and the extra refactor removes the last restart observed when running, which also reduces the total execution time from around 50s to 43s on my machine: ```console $ cd examples/bzlmod $ bazel clean --expunge --async && bazel build //entry_points:yamllint ``` Fixes bazel-contrib#1399
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -31,6 +31,12 @@ A brief description of the categories of changes:
 * (bzlmod) The `entry_point` macro is no longer supported and has been removed
   in favour of the `py_console_script_binary` macro for `bzlmod` users.
 
+### Fixed
+
+* (whl_library) No longer restarts repository rule when fetching external
+  dependencies improving initial build times involving external dependency
+  fetching.
+
 ## [0.25.0] - 2023-08-22
 
 ### Changed
diff --git a/python/pip_install/pip_repository.bzl b/python/pip_install/pip_repository.bzl
@@ -14,7 +14,7 @@
 
 ""
 
-load("//python:repositories.bzl", "get_interpreter_dirname", "is_standalone_interpreter")
+load("//python:repositories.bzl", "is_standalone_interpreter")
 load("//python:versions.bzl", "WINDOWS_NAME")
 load("//python/pip_install:repositories.bzl", "all_requirements")
 load("//python/pip_install:requirements_parser.bzl", parse_requirements = "parse")
@@ -43,15 +43,11 @@ def _construct_pypath(rctx):
     Returns: String of the PYTHONPATH.
     """
 
-    # Get the root directory of these rules
-    rules_root = rctx.path(Label("//:BUILD.bazel")).dirname
-    thirdparty_roots = [
-        # Includes all the external dependencies from repositories.bzl
-        rctx.path(Label("@" + repo + "//:BUILD.bazel")).dirname
-        for repo in all_requirements
-    ]
     separator = ":" if not "windows" in rctx.os.name.lower() else ";"
-    pypath = separator.join([str(p) for p in [rules_root] + thirdparty_roots])
+    pypath = separator.join([
+        str(rctx.path(entry).dirname)
+        for entry in rctx.attr._python_path_entries
+    ])
     return pypath
 
 def _get_python_interpreter_attr(rctx):
@@ -123,7 +119,7 @@ def _get_xcode_location_cflags(rctx):
         "-isysroot {}/SDKs/MacOSX.sdk".format(xcode_root),
     ]
 
-def _get_toolchain_unix_cflags(rctx):
+def _get_toolchain_unix_cflags(rctx, python_interpreter):
     """Gather cflags from a standalone toolchain for unix systems.
 
     Pip won't be able to compile c extensions from sdists with the pre built python distributions from indygreg
@@ -135,19 +131,19 @@ def _get_toolchain_unix_cflags(rctx):
         return []
 
     # Only update the location when using a standalone toolchain.
-    if not is_standalone_interpreter(rctx, rctx.attr.python_interpreter_target):
+    if not is_standalone_interpreter(rctx, python_interpreter):
         return []
 
     er = rctx.execute([
-        rctx.path(rctx.attr.python_interpreter_target).realpath,
+        python_interpreter,
         "-c",
         "import sys; print(f'{sys.version_info[0]}.{sys.version_info[1]}', end='')",
     ])
     if er.return_code != 0:
         fail("could not get python version from interpreter (status {}): {}".format(er.return_code, er.stderr))
     _python_version = er.stdout
     include_path = "{}/include/python{}".format(
-        get_interpreter_dirname(rctx, rctx.attr.python_interpreter_target),
+        python_interpreter.dirname,
         _python_version,
     )
 
@@ -218,19 +214,20 @@ def _parse_optional_attrs(rctx, args):
 
     return args
 
-def _create_repository_execution_environment(rctx):
+def _create_repository_execution_environment(rctx, python_interpreter):
     """Create a environment dictionary for processes we spawn with rctx.execute.
 
     Args:
-        rctx: The repository context.
+        rctx (repository_ctx): The repository context.
+        python_interpreter (path): The resolved python interpreter.
     Returns:
         Dictionary of environment variable suitable to pass to rctx.execute.
     """
 
     # Gather any available CPPFLAGS values
     cppflags = []
     cppflags.extend(_get_xcode_location_cflags(rctx))
-    cppflags.extend(_get_toolchain_unix_cflags(rctx))
+    cppflags.extend(_get_toolchain_unix_cflags(rctx, python_interpreter))
 
     env = {
         "PYTHONPATH": _construct_pypath(rctx),
@@ -630,7 +627,7 @@ def _whl_library_impl(rctx):
     result = rctx.execute(
         args,
         # Manually construct the PYTHONPATH since we cannot use the toolchain here
-        environment = _create_repository_execution_environment(rctx),
+        environment = _create_repository_execution_environment(rctx, python_interpreter),
         quiet = rctx.attr.quiet,
         timeout = rctx.attr.timeout,
     )
@@ -720,6 +717,19 @@ whl_library_attrs = {
         mandatory = True,
         doc = "Python requirement string describing the package to make available",
     ),
+    "_python_path_entries": attr.label_list(
+        # Get the root directory of these rules and keep them as a default attribute
+        # in order to avoid unnecessary repository fetching restarts.
+        #
+        # This is very similar to what was done in https://github.com/bazelbuild/rules_go/pull/3478
+        default = [
+            Label("//:BUILD.bazel"),
+        ] + [
+            # Includes all the external dependencies from repositories.bzl
+            Label("@" + repo + "//:BUILD.bazel")
+            for repo in all_requirements
+        ],
+    ),
 }
 
 whl_library_attrs.update(**common_attrs)
diff --git a/python/repositories.bzl b/python/repositories.bzl
@@ -61,39 +61,26 @@ def py_repositories():
 
 STANDALONE_INTERPRETER_FILENAME = "STANDALONE_INTERPRETER"
 
-def get_interpreter_dirname(rctx, python_interpreter_target):
-    """Get a python interpreter target dirname.
-
-    Args:
-        rctx (repository_ctx): The repository rule's context object.
-        python_interpreter_target (Target): A target representing a python interpreter.
-
-    Returns:
-        str: The Python interpreter directory.
-    """
-
-    return rctx.path(Label("{}//:WORKSPACE".format(str(python_interpreter_target).split("//")[0]))).dirname
-
-def is_standalone_interpreter(rctx, python_interpreter_target):
+def is_standalone_interpreter(rctx, python_interpreter_path):
     """Query a python interpreter target for whether or not it's a rules_rust provided toolchain
 
     Args:
         rctx (repository_ctx): The repository rule's context object.
-        python_interpreter_target (Target): A target representing a python interpreter.
+        python_interpreter_path (path): A path representing the interpreter.
 
     Returns:
         bool: Whether or not the target is from a rules_python generated toolchain.
     """
 
     # Only update the location when using a hermetic toolchain.
-    if not python_interpreter_target:
+    if not python_interpreter_path:
         return False
 
     # This is a rules_python provided toolchain.
     return rctx.execute([
         "ls",
         "{}/{}".format(
-            get_interpreter_dirname(rctx, python_interpreter_target),
+            python_interpreter_path.dirname,
             STANDALONE_INTERPRETER_FILENAME,
         ),
     ]).return_code == 0
