xangcastle/uv-1-wheel-build-memory

Author: xangcastle

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,28 @@ 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
 
+### Allocating wheel build memory
+
+Set `build_memory_mb` to the expected peak memory of a local sdist build so
+Bazel can limit concurrent wheel builds:
+
+```starlark
+uv.override_package(
+    lock = "//:uv.lock",
+    name = "native-package",
+    build_memory_mb = 6144,
+)
+```
+
+Bazel rounds the estimate up to one of its supported resource classes. On
+Linux, setting a nonzero estimate also enables aggregate resident-memory
+sampling for the build process and its descendants. Estimates above 32768 MiB
+are rejected because bazel-lib 3.2 would otherwise clamp them to 32768 MiB and
+silently under-reserve the action. 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,14 @@ 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 build resources,
+        # toolchains, and env
         # ----------------------------------------------------------------
         # Covers the override_package -> repo-rule -> BUILD-template
-        # plumbing for the augment-on-defaults `toolchains = [...]` /
+        # plumbing for `build_memory_mb`, `toolchains = [...]`, and
         # `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).
+        # default CC toolchain; the snapshot pins the memory declaration,
+        # 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,11 +23,11 @@ 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.
+# Reserve local build memory and layer a JDK runtime toolchain into the build
+# env. `toolchains` and `env` are additive to sdist_build's defaults.
 uv.override_package(
     name = "jpype1",
+    build_memory_mb = 2048,
     env = {
         "JAR": "$(JAVABASE)/bin/jar",
         "JAVA": "$(JAVA)",

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,14 @@ def make_sdist(workdir):
 
 
 def run_helper(helper, sdist, outdir, path_entries, expect):
+    rules_py_root = os.path.dirname(
+        os.path.dirname(os.path.dirname(os.path.dirname(helper)))
+    )
     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):

e2e/snapshots/sdist_build.uv_sdist_jdk_build.jpype1.BUILD.bazel

@@ -22,6 +22,7 @@ bzl_library(
         "//uv/private:normalize_version",
         "//uv/private/constraints:repository",
         "//uv/private/git_archive:repository",
+        "//uv/private/pep517_whl:build_memory",
         "//uv/private/pprint:defs",
         "//uv/private/sdist_build:repository",
         "//uv/private/sdist_configure:defs",

uv/private/extension/BUILD.bazel

@@ -60,6 +60,7 @@ load("//uv/private:normalize_version.bzl", "normalize_version")
 load("//uv/private:parse_whl_name.bzl", "parse_whl_name")
 load("//uv/private/constraints:repository.bzl", "configurations_hub")
 load("//uv/private/git_archive:repository.bzl", "git_archive")
+load("//uv/private/pep517_whl:build_memory.bzl", "validate_build_memory_mb")
 load("//uv/private/pprint:defs.bzl", "pprint")
 load("//uv/private/sdist_build:repository.bzl", "sdist_build")
 load("//uv/private/sdist_configure:defs.bzl", "DEFAULT_CONFIGURE_SCRIPT")
@@ -246,20 +247,25 @@ def _parse_projects(module_ctx, hub_specs):
                     ))
 
                 has_target = override.target != None
+                validate_build_memory_mb(
+                    override.build_memory_mb,
+                    "uv.override_package() for '{}'".format(override.name),
+                )
                 has_modifications = (
                     override.pre_build_patches or
                     override.post_install_patches or
                     override.extra_deps or
                     override.extra_data or
                     override.toolchains or
-                    override.env
+                    override.env or
+                    override.build_memory_mb
                 )
 
                 if has_target and has_modifications:
-                    fail("uv.override_package() for '{}': `target` is mutually exclusive with patch/exclude attributes. Use `target` for full replacement OR patch/exclude attributes for modifications, not both.".format(override.name))
+                    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).".format(override.name))
+                    fail("uv.override_package() for '{}': must specify either `target` for full replacement or at least one build, patch, or data modification attribute.".format(override.name))
 
                 package_overrides[override_key] = override
 
@@ -410,9 +416,11 @@ def _parse_projects(module_ctx, hub_specs):
                     # they don't replace them. Empty == no augmentation.
                     extra_toolchains = []
                     extra_env = {}
+                    build_memory_mb = 0
                     if pkg_override:
                         extra_toolchains = [str(t) for t in pkg_override.toolchains]
                         extra_env = pkg_override.env
+                        build_memory_mb = pkg_override.build_memory_mb
 
                     sbuild_specs[sbuild_id] = struct(
                         src = sdist,
@@ -425,6 +433,7 @@ def _parse_projects(module_ctx, hub_specs):
                         configure_command = project.unstable_configure_command,
                         extra_toolchains = extra_toolchains,
                         extra_env = extra_env,
+                        build_memory_mb = build_memory_mb,
                     )
 
                     has_sbuild = True
@@ -649,6 +658,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.build_memory_mb:
+            sbuild_kwargs["build_memory_mb"] = sbuild_cfg.build_memory_mb
         sdist_build(**sbuild_kwargs)
 
     for install_id, install_cfg in cfg.install_cfgs.items():
@@ -746,12 +757,16 @@ _override_package_tag = tag_class(
         # Mutually exclusive with patch/exclude attributes.
         "target": attr.label(mandatory = False),
 
-        # 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
         # env vars on top of the defaults.
+        "build_memory_mb": attr.int(
+            default = 0,
+            doc = "Estimated peak memory in MB for this package's local wheel build, from 0 to 32768. Bazel rounds up to a supported resource class; zero uses its default estimate.",
+        ),
         "toolchains": attr.label_list(
             default = [],
             doc = "Extra toolchain targets appended to the generated pep517_native_whl(...) call's `toolchains` list. Each target's TemplateVariableInfo make-variables become available for $(VAR) expansion in `env`.",
@@ -760,7 +775,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/extension/defs.bzl

@@ -1,8 +1,14 @@
 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",
+    "build_memory_failure_test",
+    "build_memory_test",
+    "pep517_native_whl_toolchain_env_test",
+)
 
 package(default_visibility = [
     "//uv/private:__subpackages__",
@@ -16,7 +22,24 @@ exports_files(
 bzl_library(
     name = "rule",
     srcs = ["rule.bzl"],
-    deps = ["//py/private/toolchain:types"],
+    deps = [
+        ":build_memory",
+        "//py/private/toolchain:types",
+        "@bazel_lib//lib:resource_sets",
+    ],
+)
+
+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(
@@ -59,6 +82,35 @@ sh_binary(
     tags = ["manual"],
 )
 
+pep517_whl(
+    name = "__excessive_memory_fixture",
+    src = ":__stub_sdist",
+    build_memory_mb = 32769,
+    tags = ["manual"],
+    tool = ":__stub_tool",
+    version = "0.0.1",
+)
+
+build_memory_failure_test(
+    name = "build_memory_limit_test",
+    expected_error = "build_memory_mb must be between 0 and 32768",
+    target_under_test = ":__excessive_memory_fixture",
+)
+
+pep517_whl(
+    name = "__maximum_memory_fixture",
+    src = ":__stub_sdist",
+    build_memory_mb = 32768,
+    tags = ["manual"],
+    tool = ":__stub_tool",
+    version = "0.0.1",
+)
+
+build_memory_test(
+    name = "build_memory_maximum_test",
+    target_under_test = ":__maximum_memory_fixture",
+)
+
 # Fixture replicates the toolchains + env defaults emitted by sdist_build's
 # BUILD template (see uv/private/sdist_build/repository.bzl). The analysis
 # test below asserts these expand to non-empty toolchain paths.
@@ -88,3 +140,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"],
+)