feat(uv): sample wheel build memory usage when resource_set is set

MODULE.bazel

@@ -6,7 +6,7 @@ module(
 )
 
 bazel_dep(name = "bazel_features", version = "1.38.0")
-bazel_dep(name = "bazel_lib", version = "3.0.0")
+bazel_dep(name = "bazel_lib", version = "3.2.0")
 bazel_dep(name = "bazel_skylib", version = "1.4.2")
 bazel_dep(name = "gawk", version = "5.3.2.bcr.7")
 bazel_dep(name = "platforms", version = "1.0.0")

docs/uv-patching.md

@@ -83,6 +83,23 @@ Pre-build patches are applied to the extracted source tree after archive extract
 - Removing problematic native build dependencies
 - Patching source code that affects the build output
 
+### Measuring wheel build memory
+
+On Linux, when `resource_set` is set, rules_py samples the aggregate resident
+memory of the build process and its descendants and reports the peak to stderr:
+
+```starlark
+uv.override_package(
+    lock = "//:uv.lock",
+    name = "native-package",
+    resource_set = "mem_2g",
+)
+```
+
+A measurement appears whenever the sampled peak crosses another 256 MiB
+threshold and once when the build finishes. If procfs is unavailable or cannot
+be read, the final measurement is reported as unavailable.
+
 ### Adding extra dependencies or data
 
 Some packages have implicit runtime dependencies that aren't declared in their metadata:

e2e/BUILD.bazel

@@ -110,13 +110,12 @@ write_source_files(
         # shows up as the per-wheel dicts collapsing into one merged list.
         "snapshots/abi3_compat.whl_install.cffi.BUILD.bazel": "@whl_install__abi3_compat__cffi__2_0_0//:BUILD.bazel",
 
-        # @sdist_build for jpype1 with uv.override_package toolchains/env
+        # @sdist_build for jpype1 with uv.override_package toolchains and env
         # ----------------------------------------------------------------
         # Covers the override_package -> repo-rule -> BUILD-template
-        # plumbing for the augment-on-defaults `toolchains = [...]` /
-        # `env = {...}`. The override layers a JDK runtime on top of the
-        # default CC toolchain; the snapshot pins both toolchain
-        # references and both env families (CC + JAVA_HOME/JAVA/JAR).
+        # plumbing for `toolchains = [...]` and `env = {...}`. The override
+        # layers a JDK runtime on top of the default CC toolchain; the snapshot
+        # pins the toolchain references and both env families (CC + JAVA_HOME/JAVA/JAR).
         # Pairs with the analysis test at
         # //uv/private/pep517_whl:toolchain_env_test which asserts the
         # rule expands env keys correctly given an explicit fixture.

e2e/cases/uv-sdist-jdk-build/setup.MODULE.bazel

@@ -23,9 +23,8 @@ uv.unstable_annotate_packages(
     lock = "//cases/uv-sdist-jdk-build:uv.lock",
 )
 
-# Layer a JDK runtime toolchain into the build env. `toolchains` and
-# `env` here are additive — these entries are appended to whatever
-# sdist_build's BUILD template emits by default.
+# Layer a JDK runtime toolchain into the build env.
+# `toolchains` and `env` are additive to sdist_build's defaults.
 uv.override_package(
     name = "jpype1",
     env = {

e2e/cases/uv-sdist-mpicc/BUILD.bazel

@@ -9,7 +9,10 @@ load("@aspect_rules_py//py:defs.bzl", "py_test")
 py_test(
     name = "test",
     srcs = ["__test__.py"],
-    data = ["@aspect_rules_py//uv/private/pep517_whl:build_helper.py"],
+    data = [
+        "@aspect_rules_py//uv/private/pep517_whl:build_helper.py",
+        "@aspect_rules_py//uv/private/pep517_whl:memory_monitor",
+    ],
     dep_group = "uv_sdist_mpicc",
     main = "__test__.py",
     python_version = "3.11",

e2e/cases/uv-sdist-mpicc/__test__.py

@@ -76,10 +76,18 @@ def make_sdist(workdir):
 
 
 def run_helper(helper, sdist, outdir, path_entries, expect):
+    from pathlib import Path as _Path
+
+    helper_path = _Path(helper)
+    assert helper_path.parts[-4:] == ("uv", "private", "pep517_whl", "build_helper.py"), (
+        "Unexpected helper path structure: {}".format(helper)
+    )
+    rules_py_root = str(helper_path.parents[3])
     env = {
         "MPICC_TEST_EXPECT": expect,
         "PATH": os.pathsep.join(path_entries),
         "HOME": os.environ.get("TEST_TMPDIR", "/tmp"),
+        "PYTHONPATH": rules_py_root,
     }
     cc = shutil.which("cc") or "/usr/bin/cc"
     if os.path.exists(cc):

uv/private/extension/defs.bzl

@@ -770,8 +770,8 @@ _override_package_tag = tag_class(
                   "bucket.",
         ),
 
-        # Per-package toolchain plumbing for native sdist builds. Both
-        # attributes AUGMENT the defaults baked into sdist_build's
+        # Per-package resource and toolchain settings for sdist builds.
+        # toolchains and env AUGMENT the defaults baked into sdist_build's
         # generated `pep517_native_whl(...)` call (the CC toolchain +
         # CC/CXX/AR/LD/STRIP env) — they don't replace them. Use these
         # to layer extra toolchains (Java runtime, Rust, …) and extra
@@ -784,7 +784,6 @@ _override_package_tag = tag_class(
             default = {},
             doc = "Extra environment variables merged into the build action's `env` dict. Values may reference $(VAR) make-variables sourced from the default CC toolchain or any extra `toolchains` listed above.",
         ),
-
         # Pre-build patches: applied to extracted sdist source before wheel build.
         "pre_build_patches": attr.label_list(
             default = [],

uv/private/pep517_whl/BUILD.bazel

@@ -1,8 +1,12 @@
 load("@bazel_lib//:bzl_library.bzl", "bzl_library")
 load("@bazel_skylib//rules:write_file.bzl", "write_file")
 load("@rules_shell//shell:sh_binary.bzl", "sh_binary")
-load(":rule.bzl", "pep517_native_whl")
-load(":test.bzl", "pep517_native_whl_toolchain_env_test")
+load("//py:defs.bzl", "py_library", "py_test")
+load(":rule.bzl", "pep517_native_whl", "pep517_whl")
+load(
+    ":test.bzl",
+    "pep517_native_whl_toolchain_env_test",
+)
 
 package(default_visibility = [
     "//uv/private:__subpackages__",
@@ -22,6 +26,19 @@ bzl_library(
     ],
 )
 
+py_library(
+    name = "memory_monitor",
+    srcs = ["memory_monitor.py"],
+    visibility = ["//visibility:public"],
+)
+
+py_test(
+    name = "memory_monitor_test",
+    srcs = ["memory_monitor_test.py"],
+    main = "memory_monitor_test.py",
+    deps = [":memory_monitor"],
+)
+
 bzl_library(
     name = "test",
     srcs = ["test.bzl"],
@@ -91,3 +108,8 @@ pep517_native_whl_toolchain_env_test(
     name = "toolchain_env_test",
     target_under_test = ":__toolchain_env_fixture",
 )
+
+bzl_library(
+    name = "build_memory",
+    srcs = ["build_memory.bzl"],
+)

uv/private/pep517_whl/build_helper.py

@@ -13,9 +13,11 @@
 import shutil
 import sys
 from os import chmod, defpath, listdir, makedirs, path, pathsep
-from subprocess import CalledProcessError, check_call, check_output, STDOUT, run
+from subprocess import CalledProcessError, check_call, check_output, STDOUT
 from tempfile import TemporaryFile
 
+from uv.private.pep517_whl.memory_monitor import run_with_memory_profile
+
 try:
     import tomllib
 except ModuleNotFoundError:
@@ -187,6 +189,7 @@ def _legacy_metadata_conflicts_with_pyproject(worktree):
 PARSER = ArgumentParser()
 PARSER.add_argument("srcarchive")
 PARSER.add_argument("outdir")
+PARSER.add_argument("--monitor-memory", action="store_true")
 PARSER.add_argument("--validate-anyarch", action="store_true")
 PARSER.add_argument("--patch-strip", type=int, default=0, help="Strip count for patch (-p)")
 PARSER.add_argument("--patch", action="append", default=[], dest="patches", help="Patch file to apply (repeatable)")
@@ -279,7 +282,14 @@ def _legacy_metadata_conflicts_with_pyproject(worktree):
 
 with TemporaryFile(mode="w+") as build_log:
     try:
-        run(cmd, cwd=t, env=build_env, stdout=build_log, stderr=STDOUT, check=True)
+        run_with_memory_profile(
+            cmd,
+            cwd=t,
+            env=build_env,
+            stdout=build_log,
+            wheel=path.basename(opts.srcarchive),
+            monitor=opts.monitor_memory,
+        )
     except CalledProcessError:
         build_log.seek(0)
         output = build_log.read()

uv/private/pep517_whl/build_memory.bzl

@@ -0,0 +1,12 @@
+"""Validation shared by wheel build configuration surfaces."""
+
+# bazel-lib 3.2 clamps larger resource requests to 32 GiB:
+# https://github.com/aspect-build/bazel-lib/blob/v3.2.0/lib/private/resource_sets.bzl#L2224-L2227
+MAX_BUILD_MEMORY_MB = 32768
+
+def validate_build_memory_mb(value, owner):
+    if value < 0 or value > MAX_BUILD_MEMORY_MB:
+        fail("{}: build_memory_mb must be between 0 and {}".format(
+            owner,
+            MAX_BUILD_MEMORY_MB,
+        ))

uv/private/pep517_whl/memory_monitor.py

@@ -0,0 +1,168 @@
+"""Run a wheel build while reporting its aggregate process-tree memory."""
+
+import os
+import signal
+import subprocess
+import sys
+import time
+
+_MIB = 1024 * 1024
+_REPORT_STEP_BYTES = 256 * _MIB
+_SAMPLE_INTERVAL_SECONDS = 0.25
+
+
+def _exit_on_sigterm(signum, _frame):
+    raise SystemExit(128 + signum)
+
+
+_children_unavailable_warned = False
+
+
+def _process_tree_rss_bytes(root_pid, page_size, proc_root="/proc"):
+    """Return aggregate RSS for root_pid and its descendants on Linux."""
+    global _children_unavailable_warned
+    # statm reports resident memory in pages:
+    # https://man7.org/linux/man-pages/man5/proc_pid_statm.5.html
+    #
+    # Each task's children file lists only immediate children, so recurse:
+    # https://docs.kernel.org/filesystems/proc.html#proc-pid-task-tid-children-information-about-task-children
+    pending = [root_pid]
+    seen = set()
+    total = 0
+    while pending:
+        pid = pending.pop()
+        if pid in seen:
+            continue
+        seen.add(pid)
+        process_dir = os.path.join(proc_root, str(pid))
+
+        try:
+            with open(os.path.join(process_dir, "statm")) as statm:
+                total += int(statm.read().split()[1]) * page_size
+            task_ids = os.listdir(os.path.join(process_dir, "task"))
+        except (
+            FileNotFoundError,
+            PermissionError,
+            ProcessLookupError,
+            ValueError,
+            IndexError,
+        ):
+            continue
+
+        any_children_file_found = False
+        for task_id in task_ids:
+            try:
+                with open(
+                    os.path.join(process_dir, "task", task_id, "children")
+                ) as children:
+                    pending.extend(int(child) for child in children.read().split())
+                any_children_file_found = True
+            except FileNotFoundError:
+                pass  # children file absent; may indicate missing CONFIG_PROC_CHILDREN
+            except (PermissionError, ProcessLookupError, ValueError):
+                any_children_file_found = True  # file exists; transient access error
+                continue
+
+        if task_ids and not any_children_file_found and not _children_unavailable_warned:
+            _children_unavailable_warned = True
+            print(
+                "rules_py: /proc/{pid}/task/{tid}/children is unavailable; "
+                "memory sampling covers only the root process "
+                "(kernel may lack CONFIG_PROC_CHILDREN).",
+                file=sys.stderr,
+                flush=True,
+            )
+
+    return total
+
+
+def _report_memory(wheel, state, current, peak):
+    if peak is None:
+        details = "unavailable"
+    elif current is None:
+        details = "sampled aggregate peak={:.1f} MiB".format(peak / _MIB)
+    else:
+        details = "sampled aggregate current={:.1f} MiB, peak={:.1f} MiB".format(
+            current / _MIB,
+            peak / _MIB,
+        )
+    print(
+        "rules_py wheel build memory for {} {}: {}".format(
+            wheel,
+            state,
+            details,
+        ),
+        file=sys.stderr,
+        flush=True,
+    )
+
+
+def run_with_memory_profile(cmd, *, cwd, env, stdout, wheel, monitor):
+    """Run cmd, optionally reporting process-tree memory."""
+    _sp_kwargs = {"cwd": cwd, "env": env, "stdout": stdout, "stderr": subprocess.STDOUT}
+
+    if not monitor:
+        subprocess.run(cmd, check=True, **_sp_kwargs)
+        return
+
+    can_sample_proc = sys.platform.startswith("linux") and os.path.isdir("/proc")
+    try:
+        page_size = os.sysconf("SC_PAGE_SIZE") if can_sample_proc else None
+    except (OSError, ValueError):
+        can_sample_proc = False
+        page_size = None
+
+    if not can_sample_proc:
+        try:
+            subprocess.run(cmd, check=True, **_sp_kwargs)
+        except subprocess.CalledProcessError:
+            _report_memory(wheel, "finished", None, None)
+            raise
+        _report_memory(wheel, "finished", None, None)
+        return
+
+    previous_sigterm_handler = signal.signal(signal.SIGTERM, _exit_on_sigterm)
+    process = None
+
+    try:
+        process = subprocess.Popen(
+            cmd,
+            start_new_session=True,
+            **_sp_kwargs,
+        )
+        peak = 0
+        next_report = _REPORT_STEP_BYTES
+
+        while True:
+            if can_sample_proc:
+                try:
+                    current = _process_tree_rss_bytes(process.pid, page_size)
+                except (OSError, ValueError):
+                    can_sample_proc = False
+                    peak = None
+                else:
+                    peak = max(peak, current)
+                    if peak >= next_report:
+                        _report_memory(wheel, "running", current, peak)
+                        next_report = (
+                            peak // _REPORT_STEP_BYTES + 1
+                        ) * _REPORT_STEP_BYTES
+
+            returncode = process.poll()
+            if returncode is not None:
+                break
+            time.sleep(_SAMPLE_INTERVAL_SECONDS)
+    except BaseException:
+        if process is not None:
+            try:
+                os.killpg(process.pid, signal.SIGKILL)
+            except ProcessLookupError:
+                pass
+            process.wait()
+        raise
+    finally:
+        signal.signal(signal.SIGTERM, previous_sigterm_handler)
+
+    _report_memory(wheel, "finished", None, peak or None)
+    if returncode:
+        raise subprocess.CalledProcessError(returncode, cmd)