feat(uv): Support native sdist builds via RBE

MODULE.bazel

@@ -27,9 +27,7 @@ use_repo(tools, "rules_py_pex_2_3_1")
 
 register_toolchains(
     "@rules_py_tools//:all",
-    "@aspect_rules_py//py/private/toolchain/venv/...",
-    "@aspect_rules_py//py/private/toolchain/unpack/...",
-    "@aspect_rules_py//py/private/toolchain/shim/...",
+    "@aspect_rules_py//py/private/toolchain/...",
 )
 
 toml = use_extension("//uv/private/tomltool:extension.bzl", "tomltool")

e2e/.bazelrc

@@ -6,4 +6,5 @@ common --@rules_cc//cc/toolchains/args/archiver_flags:use_libtool_on_macos=false
 common --repo_env=BAZEL_DO_NOT_DETECT_CPP_TOOLCHAIN=1
 common --repo_env=BAZEL_NO_APPLE_CPP_TOOLCHAIN=1
 
-
+# Configure a default venv
+common --@pypi//venv=say

e2e/MODULE.bazel

@@ -16,6 +16,7 @@ local_path_override(
     path = "..",
 )
 
+# FIXME: Can we clean this up now that sdist builds don't use tar
 bazel_lib_toolchains = use_extension("@tar.bzl//tar:extensions.bzl", "toolchains", dev_dependency = True)
 use_repo(bazel_lib_toolchains, "bsd_tar_toolchains")
 

py/private/toolchain/BUILD.bazel

@@ -1,5 +1,5 @@
 load("@bazel_lib//:bzl_library.bzl", "bzl_library")
-load(":tools.bzl", "resolved_unpack_toolchain", "resolved_venv_toolchain")
+load(":tools.bzl", "dummy_toolchain", "resolved_unpack_toolchain", "resolved_venv_toolchain")
 
 exports_files(
     ["python.sh"],
@@ -31,6 +31,26 @@ resolved_unpack_toolchain(
     visibility = ["//visibility:public"],
 )
 
+# Implementation detail of the target exec toolchain
+dummy_toolchain(
+    name = "empty",
+    visibility = ["//visibility:private"],
+)
+
+toolchain_type(
+    name = "target_exec_toolchain_type",
+)
+
+# This creates a toolchain instance with exec_compatible_with placement
+# constraints matching the target, backed by the Python toolchain already
+# resolved for the target.
+toolchain(
+    name = "target_exec_toolchain",
+    toolchain = ":empty",
+    toolchain_type = "target_exec_toolchain_type",
+    use_target_platform_constraints = True,
+)
+
 bzl_library(
     name = "autodetecting",
     srcs = ["autodetecting.bzl"],

py/private/toolchain/tools.bzl

@@ -147,3 +147,15 @@ resolved_unpack_toolchain = rule(
     implementation = _resolved_unpack_impl,
     toolchains = [UNPACK_TOOLCHAIN],
 )
+
+def _dummy_toolchain_impl(ctx):
+    toolchain_info = platform_common.ToolchainInfo(
+        dummy = True,
+    )
+    return [toolchain_info]
+
+dummy_toolchain = rule(
+    implementation = _dummy_toolchain_impl,
+    attrs = {
+    },
+)

py/private/toolchain/types.bzl

@@ -7,6 +7,7 @@ SH_TOOLCHAIN = "@bazel_tools//tools/sh:toolchain_type"
 SHIM_TOOLCHAIN = "@aspect_rules_py//py/private/toolchain:shim_toolchain_type"
 UNPACK_TOOLCHAIN = "@aspect_rules_py//py/private/toolchain:unpack_toolchain_type"
 VENV_TOOLCHAIN = "@aspect_rules_py//py/private/toolchain:venv_toolchain_type"
+TARGET_EXEC_TOOLCHAIN = "@aspect_rules_py//py/private/toolchain:target_exec_toolchain_type"
 
 PyToolInfo = provider(
     doc = "An info so we don't just return bare files",

uv/private/extension.bzl

@@ -257,10 +257,20 @@ def _parse_annotations(module_ctx, hub_specs, venv_specs):
 
     Dep = TypedDict({"name": str})
     record(
-       per_package=Dict[str, List[Dep]],
+       per_package=Dict[str, Annotation],
        default_build_deps=List[Dep],
     )
 
+    The annotations file is structured as follows:
+
+    ```
+    version = 0.0.0
+    [[package]]
+    name = "foo"
+    native = true
+    build-dependencies = [ {name = "bar"} ]
+    ```
+
     Args:
         module_ctx (module_ctx): The Bazel module context
         hub_specs (dict): The previously parsed hub specs
@@ -543,7 +553,29 @@ def _sbuild_repos(_module_ctx, lock_specs, annotation_specs, override_specs):
                     for it in build_deps + venv_anns.default_build_deps
                 }
 
-                # print("Creating sdist repo", name)
+                # TODO: Need a better native strategy.
+                #
+                # We need to decide if the sdist we're building contains native
+                # extensions. For now we're relying on user-provided annotations
+                # to do that.
+                #
+                # It would be better for the sdist_build repo rule to inspect
+                # the sdist's contents and search for well known file types such
+                # as `pyx` and `cxx`. That would also make it easier to support
+                # for instance Rust extensions. This would work from a phasing
+                # perspective because we declare sdist archives separately, so
+                # the repo rule for deciding whether an sdist build is native or
+                # not can run (prebuilt) tools to inspect downloaded archives.
+                is_native = (
+                    # Cythonized code
+                    "cython" in build_deps or
+                    # Mypyc code
+                    "mypy" in build_deps or
+                    # User has annotated the package as using C extensions or such
+                    venv_anns.per_package.get(package["name"], {}).get("native", False)
+                )
+
+                # TODO: What if the build needs toolchains? (cc, etc.)
                 sdist_build(
                     name = name,
                     src = "@" + _sdist_repo_name(package) + "//file",
@@ -552,6 +584,7 @@ def _sbuild_repos(_module_ctx, lock_specs, annotation_specs, override_specs):
                         "@" + _venv_target(hub_name, venv_name, package["name"])
                         for package in build_deps.values()
                     ],
+                    is_native = is_native,
                 )
 
 def _whl_install_repo_name(hub, venv, package):

uv/private/sdist_build/build_helper.py

@@ -1,40 +1,54 @@
 #!/usr/bin/env python3
 
+"""
+A minimal python3 -m build wrapper
+
+Mostly exists to allow debugging.
+"""
+
 from argparse import ArgumentParser
 import shutil
 import sys
-from os import getenv, listdir, path
-from subprocess import call
-
-# Under Bazel, the source dir of a sdist to build is immutable. `build` and
-# other tools however are constitutionally incapable of not writing to the
-# source tree.
-#
-# As a workaround, we use this launcher which exists to do two things.
-# - It makes a writable tempdir with a copy of the source tree
-# - It punts to `build` targeting the tempdir
-
-print(sys.executable, file=sys.stderr)
+from os import listdir, path
+from subprocess import CalledProcessError, check_call
+from tempfile import TemporaryDirectory
 
 PARSER = ArgumentParser()
-PARSER.add_argument("srcdir")
+PARSER.add_argument("srcarchive")
 PARSER.add_argument("outdir")
+PARSER.add_argument("--validate-anyarch", action="store_true")
+PARSER.add_argument("--sandbox", action="store_true")
 opts, args = PARSER.parse_known_args()
 
-t = getenv("TMPDIR")  # Provided by Bazel
-
-# Dirty awful way to prevent permissions from being replicated
-shutil.copystat = lambda x, y, **k: None
-shutil.copytree(opts.srcdir, t, dirs_exist_ok=True)
-
-outdir = path.abspath(opts.outdir)
-
-call([
-    sys.executable,
-    "-m", "build",
-    "--wheel",
-    "--no-isolation",
-    "--outdir", outdir,
-], cwd=t)
-
-print(listdir(outdir), file=sys.stderr)
+with TemporaryDirectory() as t:
+    # Extract the source archive
+    shutil.unpack_archive(opts.srcarchive, t)
+
+    # Annoyingly, unpack_archive creates a subdir in the target. Update t
+    # accordingly. Not worth the eng effort to prevent creating this dir.
+    t = path.join(t, listdir(t)[0])
+
+    # Get a path to the outdir which will be valid after we cd
+    outdir = path.abspath(opts.outdir)
+
+    try:
+        check_call([
+            sys.executable,
+            "-m", "build",
+            "--wheel",
+            "--no-isolation",
+            "--outdir", outdir,
+        ], cwd=t)
+    except CalledProcessError:
+        print("Error: Build failed!\nSee {} for the sandbox".format(t), file=sys.stderr)
+        exit(1)
+
+    inventory = listdir(outdir)
+
+    if len(inventory) > 1:
+        print("Error: Built more than one wheel!\nSee {} for the sandbox".format(t), file=sys.stderr)
+        exit(1)
+
+    if opts.validate_anyarch and not inventory[0].endswith("-none-any.whl"):
+        print("Error: Target was anyarch but built a none-any wheel!\nSee {} for the sandbox".format(t), file=sys.stderr)
+        exit(1)

uv/private/sdist_build/repository.bzl

@@ -8,29 +8,33 @@ sibling `rule.bzl` file for the implementation of `sdist_build`.
 
 def _sdist_build_impl(repository_ctx):
     repository_ctx.file("BUILD.bazel", content = """
-load("@aspect_rules_py//uv/private/sdist_build:rule.bzl", "sdist_build")
+load("@aspect_rules_py//uv/private/sdist_build:rule.bzl", "{rule}")
 load("@aspect_rules_py//py/unstable:defs.bzl", "py_venv")
 
 py_venv(
     name = "build_venv",
     deps = {deps},
 )
 
-sdist_build(
+{rule}(
     name = "whl",
     src = "{src}",
     venv = ":build_venv",
+    args = [],
     visibility = ["//visibility:public"],
 )
 """.format(
         src = repository_ctx.attr.src,
         deps = repr([str(it) for it in repository_ctx.attr.deps]),
+        # FIXME: This should probably be inferred by looking at the inventory of the sdist
+        rule = "sdist_native_build" if repository_ctx.attr.is_native else "sdist_build",
     ))
 
 sdist_build = repository_rule(
     implementation = _sdist_build_impl,
     attrs = {
         "src": attr.label(),
         "deps": attr.label_list(),
+        "is_native": attr.bool(),
     },
 )