fix(py_venv): Need venv in the exec config (#571)

The 1.5.1 release was bad for similar reasons to the 1.5.0 release --
the new `py_venv` machinery was taking a source dep on a tool which
rather than pull a prebuilt artifact created an unmet dep on
`rules_rust` to implement a source build.

This second time it was because the `venv` builder tool rather than the
interpreter `shim` wasn't depended on via a toolchain. This however is
slightly harder to fix because while the `shim` tool is used in the
target configuration strictly, the `venv` tool is used in _both_ the
target and exec configurations. The original `py_binary` machinery uses
it in target, and the new `py_venv` machinery uses it in exec to build
the venv ahead of time.

That both of these problems were present in 1.5.1 wasn't detected in
integration testing because the final batch of integration test checks
were missing the `/external/` spot check to make sure that Rust hadn't
been fetched.

This changeset does a couple things:
- (refactor) Introduce a new Info provider so that our toolchain bins
are returning an Info as they should not raw Files
- (fix) Create a new rule label which refers to the resolved file of the
`venv` toolchain, and depend on it in `exec` from `py_venv` so that we
can get a prebuilt `venv` binary but will always choose one in the
correct configuration for local usage without doing a source build
- (fix) Register the new shim toolchains in our own MODULE.bazel
- (fix) Don't revert the bzlmod dev-dep patches before running the tests
for integration
- (fix) Check that running the tests for integration didn't pull rust
- (fix) Force tests across platform boundaries to exercise that we
didn't break crossbuilding

---

### Changes are visible to end-users: yes

- Searched for relevant documentation and updated as needed: yes
- Breaking change (forces users to change their own code or config): no
- Suggested release notes appear below: no

### Test plan

- [x] New test cases added
- [x] Existing cases extended to cover previously untested behavior

Author: arrdem

.github/workflows/ci.yaml

@@ -58,12 +58,6 @@ jobs:
     defaults:
       run:
         working-directory: e2e/use_release
-    # This test references pre-built artifacts from a prior release.
-    # See py/private/toolchain/tools.bzl where this is referenced from rctx.os.environ.
-    # Will need to bump this version in the future when there are breaking changes.
-    # NB: update tools/integrity.bzl at the same time!
-    env:
-      RULES_PY_RELEASE_VERSION: 1.1.0
     steps:
       - uses: actions/checkout@v4
       - run: ./minimal_download_test.sh

MODULE.bazel

@@ -39,6 +39,7 @@ register_toolchains(
     # For manual testing: comment these out to force use of pre-built binaries.
     "@aspect_rules_py//py/private/toolchain/venv/...",
     "@aspect_rules_py//py/private/toolchain/unpack/...",
+    "@aspect_rules_py//py/private/toolchain/shim/...",
 )
 
 # To allow Rust binaries in /py/tools to be built from source

e2e/use_release/minimal_download_test.sh

@@ -2,6 +2,8 @@
 
 set -o errexit -o pipefail -o nounset
 
+set -x
+
 OS="$(uname | tr '[:upper:]' '[:lower:]')"
 ARCH="$(arch)"
 ALLOWED="rules_py_tools.${OS}_${ARCH}"
@@ -19,6 +21,9 @@ PIDFILE=$(realpath ./devserver.pid)
 
     # First we produce a release artifact matrix
     mkdir artifacts
+    # WARNING: For local testing you'll have to manually do this with
+    # --platforms for linux otherwise you'll be missing artifacts. Which is
+    # mighty annoying.
     DEST=$(realpath artifacts) bazel run //tools/release:copy_release_artifacts
 
     # We kick off a dev http server on localhost
@@ -42,6 +47,12 @@ export RULES_PY_RELEASE_URL="http://localhost:$PORT/{filename}"
     cd ../..
     # Create the .orig file, whether there's a mismatch or not
     patch -p1 --backup < .bcr/patches/*.patch
+    # Write a version to the `version.bzl` file.
+    # This emulates the version stamping git will do when it makes an archive.
+    cat <<"EOF" > tools/version.bzl
+VERSION = "999.99.9"
+IS_PRERELEASE = False
+EOF
 )
 
 OUTPUT_BASE=$(mktemp -d)
@@ -87,24 +98,49 @@ fi
 
 #############
 # Smoke test
-bazel test --test_output=streamed //...
+bazel "--output_base=$OUTPUT_BASE" test --test_output=streamed //...
 
+#############
+# Demonstrate that as configured we're fully on prebuilt toolchains even for crossbuilds
+OUTPUT_BASE=$(mktemp -d)
 (
-    cd ../..
-    rm MODULE.bazel
-    mv MODULE.bazel.orig MODULE.bazel
+  cd ../..
+
+  # Check that the configured query doesn't use Rust for anything. If we're
+  # using source toolchains, then we'll get a hit for Rust here.
+  if bazel cquery 'kind("rust_binary", deps(//py/tests/py_venv_image_layer/...))' | grep "crate_index"; then
+    >&2 echo "ERROR: we still have a rust dependency"
+    exit 1
+  fi
+
+  # Demonstrate that we can do crossbuilds with the tool
+  bazel "--output_base=$OUTPUT_BASE" build //py/tests/py_venv_image_layer/...
+
+  # TODO: Note that we can't run and pass these tests because the old py_binary
+  # implementation sees a different label for the venv tool (internal file vs
+  # external repo file) and so its image tests fail if we run them here.
 )
 
+# Note that we can't check to see if we've fetched rules_rust etc. because
+# despite being dev deps they're still visible from and fetched in the parent
+# module, even if unused.
+
 #############
 # Smoke test py_venv examples
 (
   cd ../..
-  bazel run //examples/py_venv:venv -- -c 'print("Hello, world")'
-  bazel run //examples/py_venv:internal_venv
-  bazel run --stamp //examples/py_venv:internal_venv
-  bazel run //examples/py_venv:external_venv
-  bazel run --stamp //examples/py_venv:external_venv
+  # Exercise the static venv bits
+  # Note that we only really expect
+  bazel "--output_base=$OUTPUT_BASE" run //examples/py_venv:venv -- -c 'print("Hello, world")'
+  bazel "--output_base=$OUTPUT_BASE" run //examples/py_venv:internal_venv
+  bazel "--output_base=$OUTPUT_BASE" run --stamp //examples/py_venv:internal_venv
+  bazel "--output_base=$OUTPUT_BASE" run //examples/py_venv:external_venv
+  bazel "--output_base=$OUTPUT_BASE" run --stamp //examples/py_venv:external_venv
 )
 
+# Note that we can't check to see if we've fetched rules_rust etc. because
+# despite being dev deps they're still visible from and fetched in the parent
+# module, even if unused.
+
 # Shut down the devserver
 kill "$(cat $PIDFILE)"

py/extensions.bzl

@@ -31,6 +31,7 @@ def _toolchains_extension_impl(module_ctx):
                 root_name = toolchain.name
             else:
                 registrations.append(toolchain.name)
+
     for name in registrations:
         if name != root_name:
             rules_py_toolchains(name, register = False)

py/private/py_binary.bzl

@@ -73,7 +73,7 @@ def _py_binary_rule_impl(ctx):
         substitutions = {
             "{{BASH_RLOCATION_FN}}": BASH_RLOCATION_FUNCTION,
             "{{INTERPRETER_FLAGS}}": " ".join(py_toolchain.flags + ctx.attr.interpreter_options),
-            "{{VENV_TOOL}}": to_rlocation_path(ctx, venv_toolchain.bin),
+            "{{VENV_TOOL}}": to_rlocation_path(ctx, venv_toolchain.bin.bin),
             "{{ARG_COLLISION_STRATEGY}}": ctx.attr.package_collisions,
             "{{ARG_PYTHON}}": to_rlocation_path(ctx, py_toolchain.python) if py_toolchain.runfiles_interpreter else py_toolchain.python.path,
             "{{ARG_VENV_NAME}}": ".{}.venv".format(ctx.attr.name),

py/private/py_unpacked_wheel.bzl

@@ -29,7 +29,7 @@ def _py_unpacked_wheel_impl(ctx):
     ctx.actions.run(
         outputs = [unpack_directory],
         inputs = depset([ctx.file.src], transitive = [py_toolchain.files]),
-        executable = unpack_toolchain.bin,
+        executable = unpack_toolchain.bin.bin,
         arguments = [arguments],
         mnemonic = "PyUnpackedWheel",
         progress_message = "Unpacking wheel {}".format(ctx.file.src.basename),

py/private/py_venv/py_venv.bzl

@@ -5,7 +5,7 @@ load("@aspect_bazel_lib//lib:paths.bzl", "BASH_RLOCATION_FUNCTION", "to_rlocatio
 load("//py/private:py_library.bzl", _py_library = "py_library_utils")
 load("//py/private:py_semantics.bzl", _py_semantics = "semantics")
 load("//py/private:transitions.bzl", "python_version_transition")
-load("//py/private/toolchain:types.bzl", "PY_TOOLCHAIN", "SHIM_TOOLCHAIN")
+load("//py/private/toolchain:types.bzl", "PY_TOOLCHAIN", "PyToolInfo", "SHIM_TOOLCHAIN", "VENV_TOOLCHAIN")
 
 VirtualenvInfo = provider(
     doc = """
@@ -74,7 +74,10 @@ def _py_venv_base_impl(ctx):
 
     py_toolchain = _py_semantics.resolve_toolchain(ctx)
 
+    # Note that we HAVE to grab these files from toolchains so that we can swap
+    # in prebuild versions in production.
     py_shim = ctx.toolchains[SHIM_TOOLCHAIN]
+    venv_tool = ctx.attr._venv[PyToolInfo].bin
 
     # Check for duplicate virtual dependency names. Those that map to the same resolution target would have been merged by the depset for us.
     virtual_resolution = _py_library.resolve_virtuals(ctx)
@@ -141,10 +144,10 @@ def _py_venv_base_impl(ctx):
     venv_dir = ctx.actions.declare_directory(venv_name)
 
     ctx.actions.run(
-        executable = ctx.file._venv_tool,
+        executable = venv_tool,
         arguments = [
             "--location=" + venv_dir.path,
-            "--python=" + py_shim.bin.path,
+            "--python=" + py_shim.bin.bin.path,
             "--pth-file=" + site_packages_pth_file.path,
             "--env-file=" + env_file.path,
             "--bin-dir=" + ctx.bin_dir.path,
@@ -160,7 +163,7 @@ def _py_venv_base_impl(ctx):
             ctx.runfiles(files = [
                 site_packages_pth_file,
                 env_file,
-                ctx.file._venv_tool,
+                venv_tool,
             ]),
             py_shim.default_info.default_runfiles,
         ]).files,
@@ -313,12 +316,9 @@ A collision can occur when multiple packages providing the same file are install
     "_runfiles_lib": attr.label(
         default = "@bazel_tools//tools/bash/runfiles",
     ),
-    # Note that we're using the transitioned one for local execution, since
-    # we're going to run the tool on the local platform and produce static files
-    # not including this tool.
-    "_venv_tool": attr.label(
-        allow_single_file = True,
-        default = "//py/tools/venv_bin:local_venv_bin",
+    "_venv": attr.label(
+        default = "//py/private/toolchain:resolved_venv_toolchain",
+        cfg = "exec",
     ),
 })
 
@@ -365,6 +365,7 @@ py_venv_base = struct(
     toolchains = [
         PY_TOOLCHAIN,
         SHIM_TOOLCHAIN,
+        VENV_TOOLCHAIN,
     ],
     cfg = python_version_transition,
 )

py/private/toolchain/BUILD.bazel

@@ -1,4 +1,5 @@
 load("@bazel_skylib//:bzl_library.bzl", "bzl_library")
+load(":tools.bzl", "resolved_venv_toolchain")
 
 exports_files(
     ["python.sh"],
@@ -20,6 +21,11 @@ toolchain_type(
     visibility = ["//visibility:public"],
 )
 
+resolved_venv_toolchain(
+    name = "resolved_venv_toolchain",
+    visibility = ["//visibility:public"],
+)
+
 bzl_library(
     name = "autodetecting",
     srcs = ["autodetecting.bzl"],

py/private/toolchain/repo.bzl

@@ -34,13 +34,14 @@ def _toolchains_repo_impl(repository_ctx):
             build_content += """
 # Declare a toolchain Bazel will select for running {tool} on the {cfg} platform.
 toolchain(
-    name = "{tool}_{platform}_toolchain",
+    name = "{tool}_{platform}_{cfg}_toolchain",
     {cfg}_compatible_with = {compatible_with},
     # Bazel does not follow this attribute during analysis, so the referenced repo
     # will only be fetched if this toolchain is selected.
     toolchain = "@{user_repository_name}.{platform}//:{tool}_toolchain",
     toolchain_type = "{toolchain_type}",
 )
+
 """.format(
                 cfg = bin.cfg,
                 tool = bin.name,

py/private/toolchain/tools.bzl

@@ -1,5 +1,8 @@
 """Declaration of concrete toolchains for our Rust tools"""
 
+load("@bazel_skylib//lib:structs.bzl", "structs")
+load(":types.bzl", "PyToolInfo", "VENV_TOOLCHAIN")
+
 def PrebuiltToolConfig(
         target,
         cfg = "target",
@@ -82,7 +85,7 @@ def _toolchain_impl(ctx):
     # Export all the providers inside our ToolchainInfo
     # so the resolved_toolchain rule can grab and re-export them.
     toolchain_info = platform_common.ToolchainInfo(
-        bin = binary,
+        bin = PyToolInfo(bin = binary),
         template_variables = template_variables,
         default_info = default_info,
     )
@@ -122,3 +125,15 @@ def source_toolchain(name, toolchain_type, bin):
         toolchain = toolchain_rule,
         toolchain_type = toolchain_type,
     )
+
+# Forward all the providers
+def _resolved_toolchain_impl(ctx):
+    toolchain_info = ctx.toolchains[VENV_TOOLCHAIN]
+    return [toolchain_info] + structs.to_dict(toolchain_info).values()
+
+# Copied from java_toolchain_alias
+# https://cs.opensource.google/bazel/bazel/+/master:tools/jdk/java_toolchain_alias.bzl
+resolved_venv_toolchain = rule(
+    implementation = _resolved_toolchain_impl,
+    toolchains = [VENV_TOOLCHAIN],
+)