Seal Python with a VirtualEnv toolchain

BUILD

load("@rules_python//python:defs.bzl", "py_runtime", "py_runtime_pair")
load(":venv.bzl", "py_venv")
load(":wrapper.bzl", "py_wrapper")

py_venv(
    name = "py3_venv",
    # caching the venv interpreter doesn't work for two reasons:
    #
    # * there are no platform dependencies on this target: the cache will contain binaries for the wrong OS
    # * there are no dependencies on the interpreter binary or standard library: the cache may contain stale versions
    #
    # fixing this properly a bit of a pain...
    # the preferred solution is to use a platform specific tarball that does not need to be relocated.
    # unfortunately the Python organization only builds/hosts _installers_ and does not provide binary tarballs.
    # building/maintaining our own binary tarballs is quite a bit of work, if we wait someone else may do it for us.
    # the second best would to add the interpreter and standard library as inputs to the py_venv rule.
    # this is not a terrible interim solution, it just needs a repository_rule and a python script to hash
    # said interpreter and libraries, but it's not free and does not provide hermeticicity (binary tarballs would).
    #
    # "no-remote-cache" fixes bullet 1. "no-cache" fixes bullet 1 and mostly avoids bullet 2.
    tags = ["no-cache"],
)

py_wrapper(
    name = "py3_wrapper",
    venv = ":py3_venv",
)

# python2 doesn't have proper venv support so invoke it with -S and hope for the best
py_runtime(
    name = "py2_runtime",
    interpreter = ":py2wrapper.sh",
    python_version = "PY2",
)

# python3 gets a real venv, this is required to run tensorflow due to expectations that
# site.getsitepackages() returns a list of directories
py_runtime(
    name = "py3_runtime",
    files = [":py3_venv"],  # annoying: https://github.com/bazelbuild/bazel/issues/4286
    interpreter = ":py3_wrapper",
    python_version = "PY3",
)

py_runtime_pair(
    name = "py23_runtime_pair",
    py2_runtime = ":py2_runtime",
    py3_runtime = ":py3_runtime",
)

toolchain(
    name = "python",
    toolchain = ":py23_runtime_pair",
    toolchain_type = "@bazel_tools//tools/python:toolchain_type",
)

create_venv.sh

#!/usr/bin/env sh
set -e
python3.7 -m venv --without-pip "$1"

# remove unused activation scripts that break caching
rm "$1/bin/"activate*

LICENSE

Copyright (c) 2022 Nathan Howell

Permission is hereby granted, free of charge, to any person obtaining a copy
of this software and associated documentation files (the "Software"), to deal
in the Software without restriction, including without limitation the rights
to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
copies of the Software, and to permit persons to whom the Software is
furnished to do so, subject to the following conditions:

The above copyright notice and this permission notice shall be included in all
copies or substantial portions of the Software.

THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
SOFTWARE.

py3wrapper.sh

#!/usr/bin/env bash
ROOT=$(dirname "$0")
# disable pycache which produces non-deterministic outputs
export PYTHONDONTWRITEBYTECODE=1
exec "$ROOT/venv/bin/python3.7" -B "$@"

venv.bzl

def _impl(ctx):
    version = ctx.attr.version.split(".", 2)
    if len(version) != 2:
        fail("Expected version in major.minor format, e.g. 3.7. Found: {}".format(ctx.attr.version))

    py_minor = ctx.actions.declare_file("venv/bin/python{}".format(ctx.attr.version))
    py_major = ctx.actions.declare_file("venv/bin/python{}".format(version[0]))
    python = ctx.actions.declare_file("venv/bin/python")
    venv_cfg = ctx.actions.declare_file("venv/pyvenv.cfg")

    ctx.actions.run(
        executable = ctx.executable._create_venv,
        arguments = [venv_cfg.dirname],
        outputs = [python, py_major, py_minor, venv_cfg],
        use_default_shell_env = True,
    )

    return [
        DefaultInfo(
            executable = py_minor,
            runfiles = ctx.runfiles(
                files = [python, py_major, py_minor, venv_cfg],
            ),
        ),
    ]

py_venv = rule(
    implementation = _impl,
    executable = True,
    attrs = {
        "version": attr.string(default = "3.7"),
        "_create_venv": attr.label(
            default = ":create_venv.sh",
            allow_single_file = True,
            executable = True,
            cfg = "target",
        ),
    },
)

WORKSPACE

# register local toolchains before loading other workspaces
register_toolchains("//toolchains:python")

wrapper.bzl

def _impl(ctx):
    wrapper = ctx.actions.declare_file("py3wrapper.sh")

    # we are required to create a new file when returning an executable...
    # so even though no substitutions are being made this seems about the same
    # as running `cp`
    ctx.actions.expand_template(
        template = ctx.file._template,
        output = wrapper,
        is_executable = True,
        substitutions = {},
    )

    return [
        DefaultInfo(
            executable = wrapper,
            runfiles = ctx.runfiles(
                files = [wrapper],
                transitive_files = depset(ctx.files.venv),
            ),
        ),
    ]

py_wrapper = rule(
    implementation = _impl,
    executable = True,
    attrs = {
        "venv": attr.label(
            mandatory = True,
            executable = True,
            cfg = "target",
        ),
        "_template": attr.label(
            default = ":py3wrapper.sh",
            allow_single_file = True,
        ),
    },
)

Looks like this could be useful. Would you mind releasing it under a suitably permissive license (e.g. MIT License, Apache 2.0 License, or something suitable)?

@vonschultz license applied, one of these days I'll add the toolchain to the poetry rules

Thanks. I've now tried it out. A few comments:

1. In create_venv.sh:

# remove unused activation scripts that break caching
rm "$1/bin/"activate*

This fails to remove Activate.ps1, because Linux is case sensitive. To remove all activation scripts you could do [Aa]ctivate*.

2. Though the py_venv has a version attribute, it only actually supports Python 3.7, since that's hardcoded in create_venv.sh. An easy fix would be to do

"$1" -m venv --without-pip "$2"

in create_venv.sh and do

        arguments = ["python{}".format(ctx.attr.version), venv_cfg.dirname],

in venv.bazel.

3. Similarly py3wrapper.sh shouldn't hardcode $ROOT/venv/bin/python3.7. A simple solution would be to use $ROOT/venv/bin/python3 instead, since all versions of Python 3 will have this, and "py3" is in the name of the wrapper, so that assumption is fine. To support this I also needed to do executable = py_major, instead of executable = py_minor, in the implementation of py_venv.
4. The tags = ["no-cache"], in the "py3_venv" doesn't take effect in my testing, because the rule implementation ignores the tags. You would need to set the execution_requirements in the ctx.actions.run call to take tags into account. Something like execution_requirements = { tag: "true" for tag in ctx.attr.tags }
5. There is a reference to :py2wrapper.sh, but that file is not included here.

One more update:

6. I've also had to add stub_shebang = "#!/usr/bin/env python3", to the :py3_runtime, in order to support systems with no python executable, which do have python3 (e.g. Ubuntu default installation).
7. I've added files = depset([python, py_major, py_minor, venv_cfg]), to the DefaultInfo of the py_venv implementation. Not entirely sure why that was needed, since it already has runfiles, but some things refused to work without it.

@vonschultz thanks for this, as you noticed it was implemented just enough to work for our application and never was buttoned up. if you'd like to fork the gist and push updates I can merge them in.

OK, I've pushed updates to a fork, covering most of my points. I don't have any :py2wrapper.sh, though.