feat(uv): monitor wheel build memory (#1174)

Large source-built wheels can exhaust memory before their action logs
show
enough evidence to size a scheduler reservation.

Add an explicit `monitor_memory` package override that samples
approximate
Linux procfs RSS for the PEP 517 frontend and its descendants. Emit the
first sample, 256 MiB high-water crossings, and a final peak. Preserve
earlier samples across procfs races and identify a direct SIGKILL as a
possible OOM.

Keep diagnostics independent from `resource_set`. The monitor inherits
the
Bazel process group and never installs signal handlers or kills child
processes, so Bazel retains cancellation ownership. Only opted-in
generated
build tools carry the monitor dependency.

Apply monitoring only to generated sdist actions. Selecting a compatible
prebuilt wheel produces no report; prebuilt-only packages and complete
custom BUILD replacements reject the override.

Author: tamird

docs/uv-patching.md

@@ -126,6 +126,40 @@ package that resolves to a prebuilt wheel (no source build) fails the build
 rather than silently dropping the reservation — force a source build with
 `[tool.uv] no-binary-package` if you need the reservation to apply.
 
+### Monitoring wheel build memory
+
+Set `monitor_memory` to report the memory observed while building a wheel from
+an sdist:
+
+```starlark
+uv.override_package(
+    lock = "//:uv.lock",
+    name = "native-package",
+    monitor_memory = True,
+)
+```
+
+On Linux, rules_py reports the first sample, each 256 MiB high-water crossing,
+and the final peak. Reports are flushed as the build runs, so an earlier
+high-water mark can remain in the action log when an OOM kills the build.
+
+The measurement is a best-effort sum of `/proc` RSS for the build process and
+its descendants. It can double-count shared pages and miss short-lived
+processes. On other platforms it is reported as unavailable.
+
+`monitor_memory` is diagnostic only. It neither limits memory nor reserves
+scheduler capacity, and can be enabled independently from `resource_set`.
+
+Monitoring runs only when the source-build target is selected. A package with
+both an sdist and a compatible wheel produces no report when the wheel is
+selected; use `[tool.uv] no-binary-package` to force the monitored source build.
+A package with no sdist rejects the override.
+
+A custom sdist configure tool that returns complete `build_file_content` owns
+the wheel action itself. Such content must add its own monitoring; combining
+that replacement with `monitor_memory` is rejected rather than silently
+dropping the diagnostic.
+
 ### Full replacement
 
 To replace a package entirely with a custom target (existing functionality):

e2e/BUILD.bazel

@@ -128,6 +128,14 @@ write_source_files(
         # rule expands env keys correctly given an explicit fixture.
         "snapshots/sdist_build.uv_sdist_jdk_build.jpype1.BUILD.bazel": "@sdist_build__uv_sdist_jdk_build__jpype1__1_7_1//:BUILD.bazel",
 
+        # @sdist_build for cowsay with uv.override_package(monitor_memory)
+        # ----------------------------------------------------------------
+        # Covers the override_package -> repo-rule -> BUILD-template plumbing
+        # for a monitor-only source build. The runtime
+        # //cases/uv-sdist-fallback:test proves that the generated build tool
+        # can load the conditional monitor dependency and complete the wheel.
+        "snapshots/sdist_build.uv_sdist_fallback.cowsay.BUILD.bazel": "@sdist_build__uv_sdist_fallback__cowsay__6_0//:BUILD.bazel",
+
         # @sdist_build for python-geohash with uv.override_package(resource_set)
         # ----------------------------------------------------------------------
         # Covers the override_package -> repo-rule -> BUILD-template plumbing

e2e/cases/uv-sdist-fallback/setup.MODULE.bazel

@@ -16,4 +16,13 @@ uv.project(
     lock = "//cases/uv-sdist-fallback:uv.lock",
     pyproject = "//cases/uv-sdist-fallback:pyproject.toml",
 )
-use_repo(uv, "pypi_sdist_fallback")
+uv.override_package(
+    name = "cowsay",
+    lock = "//cases/uv-sdist-fallback:uv.lock",
+    monitor_memory = True,
+)
+use_repo(
+    uv,
+    "pypi_sdist_fallback",
+    "sdist_build__uv_sdist_fallback__cowsay__6_0",
+)

e2e/snapshots/sdist_build.uv_sdist_fallback.cowsay.BUILD.bazel

@@ -254,14 +254,15 @@ def _parse_projects(module_ctx, hub_specs):
                     override.extra_data or
                     override.toolchains or
                     override.env or
+                    override.monitor_memory or
                     override.resource_set != "default"
                 )
 
                 if has_target and has_modifications:
                     fail("uv.override_package() for '{}': `target` is mutually exclusive with modification attributes. Use `target` for full replacement OR build, patch, and data attributes for modifications, not both.".format(override.name))
 
                 if not has_target and not has_modifications:
-                    fail("uv.override_package() for '{}': must specify either `target` for full replacement or at least one modification attribute (pre_build_patches, post_install_patches, extra_deps, extra_data, toolchains, env, resource_set).".format(override.name))
+                    fail("uv.override_package() for '{}': must specify either `target` for full replacement or at least one modification attribute (pre_build_patches, post_install_patches, extra_deps, extra_data, toolchains, env, monitor_memory, resource_set).".format(override.name))
 
                 package_overrides[override_key] = override
 
@@ -412,10 +413,12 @@ def _parse_projects(module_ctx, hub_specs):
                     # they don't replace them. Empty == no augmentation.
                     extra_toolchains = []
                     extra_env = {}
+                    monitor_memory = False
                     resource_set = "default"
                     if pkg_override:
                         extra_toolchains = [str(t) for t in pkg_override.toolchains]
                         extra_env = pkg_override.env
+                        monitor_memory = pkg_override.monitor_memory
                         resource_set = pkg_override.resource_set
 
                     sbuild_specs[sbuild_id] = struct(
@@ -429,6 +432,7 @@ def _parse_projects(module_ctx, hub_specs):
                         configure_command = project.unstable_configure_command,
                         extra_toolchains = extra_toolchains,
                         extra_env = extra_env,
+                        monitor_memory = monitor_memory,
                         resource_set = resource_set,
                     )
 
@@ -448,6 +452,8 @@ def _parse_projects(module_ctx, hub_specs):
 
                 if pkg_override and pkg_override.resource_set != "default" and not has_sbuild:
                     fail("uv.override_package() for '{}': `resource_set` reserves resources for the sdist wheel-build action, but this package resolves to a prebuilt wheel (there is no sdist build to reserve for). Remove `resource_set`, or force a source build via `[tool.uv] no-binary-package`.".format(pkg_override.name))
+                if pkg_override and pkg_override.monitor_memory and not has_sbuild:
+                    fail("uv.override_package() for '{}': `monitor_memory` observes the sdist wheel-build action, but this package resolves to a prebuilt wheel. Remove `monitor_memory`, or force a source build via `[tool.uv] no-binary-package`.".format(pkg_override.name))
 
                 # uv can emit multiple lock records for the same package/version
                 # (e.g. resolution-marker forks), each carrying a different
@@ -657,6 +663,8 @@ def _uv_impl(module_ctx):
             sbuild_kwargs["extra_toolchains"] = sbuild_cfg.extra_toolchains
         if sbuild_cfg.extra_env:
             sbuild_kwargs["extra_env"] = sbuild_cfg.extra_env
+        if sbuild_cfg.monitor_memory:
+            sbuild_kwargs["monitor_memory"] = True
         if sbuild_cfg.resource_set != "default":
             sbuild_kwargs["resource_set"] = sbuild_cfg.resource_set
         sdist_build(**sbuild_kwargs)
@@ -755,6 +763,10 @@ _override_package_tag = tag_class(
         # Full replacement: provide a target that substitutes for the package entirely.
         # Mutually exclusive with patch/exclude attributes.
         "target": attr.label(mandatory = False),
+        "monitor_memory": attr.bool(
+            default = False,
+            doc = "Report approximate Linux process-tree RSS while building this package's wheel.",
+        ),
 
         # Per-package local execution resources for the wheel build action.
         # Uses bazel-lib's predefined resource_set vocabulary (the same enum

uv/private/extension/defs.bzl

@@ -4,7 +4,7 @@ load("@bazel_skylib//rules:write_file.bzl", "write_file")
 load("@rules_shell//shell:sh_binary.bzl", "sh_binary")
 load("@tar.bzl//tar:mtree.bzl", "mtree_mutate", "mtree_spec")
 load("@tar.bzl//tar:tar.bzl", "tar_rule")
-load("//py:defs.bzl", "py_binary")
+load("//py:defs.bzl", "py_binary", "py_library", "py_test")
 load(":rule.bzl", "pep517_native_whl", "pep517_whl")
 load(":test.bzl", "hostile_python_env_target", "pep517_native_whl_toolchain_env_test")
 
@@ -26,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"],
+)
+
 py_binary(
     name = "__build_helper",
     testonly = True,

uv/private/pep517_whl/BUILD.bazel

@@ -187,6 +187,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 +280,20 @@ 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)
+        if opts.monitor_memory:
+            # Generated build tools include this dependency only when the
+            # corresponding wheel opts into monitoring.
+            from uv.private.pep517_whl.memory_monitor import run_with_memory_monitor
+
+            run_with_memory_monitor(
+                cmd,
+                cwd=t,
+                env=build_env,
+                stdout=build_log,
+                wheel=path.basename(opts.srcarchive),
+            )
+        else:
+            run(cmd, cwd=t, env=build_env, stdout=build_log, stderr=STDOUT, check=True)
     except CalledProcessError:
         build_log.seek(0)
         output = build_log.read()

uv/private/pep517_whl/build_helper.py

@@ -0,0 +1,116 @@
+"""Report approximate Linux RSS while a wheel build runs."""
+
+import os
+import subprocess
+import sys
+import time
+
+_MIB = 1024 * 1024
+_REPORT_STEP_BYTES = 256 * _MIB
+_SAMPLE_INTERVAL_SECONDS = 0.25
+
+
+def _process_tree_rss_bytes(root_pid, page_size, proc_root="/proc"):
+    """Return the sampled RSS sum for root_pid and its descendants."""
+    # statm reports resident memory in pages:
+    # https://man7.org/linux/man-pages/man5/proc_pid_statm.5.html
+    # Each task's children file lists that task's immediate children:
+    # https://docs.kernel.org/filesystems/proc.html#proc-pid-task-tid-children-information-about-task-children
+    pending = [root_pid]
+    seen = set()
+    total = 0
+    sampled = False
+
+    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
+            sampled = True
+        except (FileNotFoundError, PermissionError, ProcessLookupError, ValueError, IndexError):
+            pass
+
+        try:
+            task_ids = os.listdir(os.path.join(process_dir, "task"))
+        except (FileNotFoundError, PermissionError, ProcessLookupError):
+            continue
+
+        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())
+            except (FileNotFoundError, PermissionError, ProcessLookupError, ValueError):
+                continue
+
+    return total if sampled else None
+
+
+def _report_memory(wheel, state, peak, current=None, returncode=None):
+    if peak is None:
+        details = "unavailable"
+    elif current is None:
+        details = "peak={:.1f} MiB".format(peak / _MIB)
+    else:
+        details = "current={:.1f} MiB, peak={:.1f} MiB".format(
+            current / _MIB,
+            peak / _MIB,
+        )
+    if returncode is not None:
+        details += ", return code={}".format(returncode)
+        if returncode == -9:
+            details += " (SIGKILL; possible OOM)"
+    print(
+        "rules_py wheel build memory for {} {}: {} (approximate aggregate RSS)".format(
+            wheel,
+            state,
+            details,
+        ),
+        file=sys.stderr,
+        flush=True,
+    )
+
+
+def run_with_memory_monitor(cmd, *, cwd, env, stdout, wheel):
+    """Run cmd while reporting best-effort process-tree RSS."""
+    try:
+        can_sample = sys.platform.startswith("linux") and os.path.isdir("/proc")
+        page_size = os.sysconf("SC_PAGE_SIZE") if can_sample else None
+    except (OSError, ValueError):
+        can_sample = False
+        page_size = None
+
+    process = subprocess.Popen(
+        cmd,
+        cwd=cwd,
+        env=env,
+        stdout=stdout,
+        stderr=subprocess.STDOUT,
+    )
+    peak = None
+    next_report = 0
+
+    while True:
+        if can_sample:
+            try:
+                current = _process_tree_rss_bytes(process.pid, page_size)
+            except (OSError, ValueError):
+                current = None
+            if current is not None:
+                peak = current if peak is None else max(peak, current)
+                if peak >= next_report:
+                    _report_memory(wheel, "running", peak, current=current)
+                    next_report = (peak // _REPORT_STEP_BYTES + 1) * _REPORT_STEP_BYTES
+
+        returncode = process.poll()
+        if returncode is not None:
+            break
+        time.sleep(_SAMPLE_INTERVAL_SECONDS)
+
+    _report_memory(wheel, "finished", peak, returncode=returncode)
+    if returncode:
+        raise subprocess.CalledProcessError(returncode, cmd)