From c47529435773ea6620a82037f20c0482e8c77803 Mon Sep 17 00:00:00 2001 From: Alex Eagle Date: Fri, 12 Jul 2024 09:49:51 -0700 Subject: [PATCH] feat: provide affordance for OCIImage actions to reserve more resources (#650) --- docs/image.md | 3 +- examples/big_image/BUILD.bazel | 19 +++----- oci/private/BUILD.bazel | 11 ++++- oci/private/image.bzl | 6 ++- oci/private/resource_sets.bzl | 89 ++++++++++++++++++++++++++++++++++ 5 files changed, 113 insertions(+), 15 deletions(-) create mode 100644 oci/private/resource_sets.bzl diff --git a/docs/image.md b/docs/image.md index 3b830bc0..f1bf426f 100644 --- a/docs/image.md +++ b/docs/image.md @@ -14,7 +14,7 @@ load("@rules_oci//oci:defs.bzl", ...) 
 oci_image_rule(name, annotations, architecture, base, cmd, entrypoint, env, exposed_ports, labels,
-               os, tars, user, variant, workdir)
+               os, resource_set, tars, user, variant, workdir)
Build an OCI compatible container image. @@ -80,6 +80,7 @@ oci_image( | exposed_ports | A file containing a comma separated list of exposed ports. (e.g. 2000/tcp, 3000/udp or 4000. No protocol defaults to tcp). | Label | optional | None | | labels | A file containing a dictionary of labels. Each line should be in the form name=value. | Label | optional | None | | os | The name of the operating system which the image is built to run on. eg: linux, windows. See $GOOS documentation for possible values: https://go.dev/doc/install/source#environment | String | optional | "" | +| resource_set | A predefined function used as the resource_set for actions.

Used with --experimental_action_resource_set to reserve more RAM/CPU, preventing Bazel overscheduling resource-intensive actions.

By default, Bazel allocates 1 CPU and 250M of RAM. https://github.com/bazelbuild/bazel/blob/058f943037e21710837eda9ca2f85b5f8538c8c5/src/main/java/com/google/devtools/build/lib/actions/AbstractAction.java#L77 | String | optional | "default" | | tars | List of tar files to add to the image as layers. Do not sort this list; the order is preserved in the resulting image. Less-frequently changed files belong in lower layers to reduce the network bandwidth required to pull and push.

The authors recommend [dive](https://github.com/wagoodman/dive) to explore the layering of the resulting image. | List of labels | optional | [] | | user | The username or UID which is a platform-specific structure that allows specific control over which user the process run as. This acts as a default value to use when the value is not specified when creating a container. For Linux based systems, all of the following are valid: user, uid, user:group, uid:gid, uid:group, user:gid. If group/gid is not specified, the default group and supplementary groups of the given user/uid in /etc/passwd from the container are applied. | String | optional | "" | | variant | The variant of the specified CPU architecture. eg: v6, v7, v8. See: https://github.com/opencontainers/image-spec/blob/main/image-index.md#platform-variants for more. | String | optional | "" | diff --git a/examples/big_image/BUILD.bazel b/examples/big_image/BUILD.bazel index e916a464..d01c97d6 100644 --- a/examples/big_image/BUILD.bazel +++ b/examples/big_image/BUILD.bazel @@ -1,19 +1,18 @@ load("@rules_pkg//pkg:tar.bzl", "pkg_tar") load("//oci:defs.bzl", "oci_image") - # These numbers were gathered on a `Apple M2 Pro` # Darwin Kernel Version 23.2.0: Wed Nov 15 21:55:06 PST 2023; root:xnu-10002.61.3~2/RELEASE_ARM64_T6020 # 10 CPU x 32GB RAM # 1- Create an image with 10 layers 3GiB each. # Perf: `< 50s` -N_BASE_LAYERS=10 +N_BASE_LAYERS = 10 genrule( name = "3gib_file", - cmd = "mkfile 3G $@", outs = ["3gib_file.out"], + cmd = "mkfile 3G $@", tags = ["manual"], ) @@ -29,10 +28,10 @@ genrule( oci_image( name = "base", - os = "linux", architecture = "arm64", - tars = [":blayer_%s" % i for i in range(N_BASE_LAYERS)], + os = "linux", tags = ["manual"], + tars = [":blayer_%s" % i for i in range(N_BASE_LAYERS)], ) # 2- Create an image that extends the base without adding additional layers @@ -44,13 +43,11 @@ oci_image( tags = ["manual"], ) - - # 3- Create an image that extends the base with additional layers # Perf: `< 2s` if bazel doesn't permit tree artifact symlinks (for copying the base layers) # Perf: `< 0.5s` if bazel permits tree artifact symlinks (for linking the base layers) # Perf: `< 20s` for adding the new 5 layers, 3gib each. -N_LAYERS=5 +N_LAYERS = 5 [ pkg_tar( @@ -61,11 +58,11 @@ N_LAYERS=5 for i in range(N_LAYERS) ] - oci_image( name = "extended", base = ":base", - tars = [":layer_%s" % i for i in range(N_LAYERS)], + # Tell Bazel to reserve more than the default 250MB of RAM for the OCIImage action + resource_set = "mem_8g", tags = ["manual"], + tars = [":layer_%s" % i for i in range(N_LAYERS)], ) - diff --git a/oci/private/BUILD.bazel b/oci/private/BUILD.bazel index a82e3d52..7ec9081c 100644 --- a/oci/private/BUILD.bazel +++ b/oci/private/BUILD.bazel @@ -27,7 +27,10 @@ bzl_library( "//docs:__pkg__", "//oci:__subpackages__", ], - deps = [":util"], + deps = [ + ":resource_sets", + ":util", + ], ) bzl_library( @@ -108,3 +111,9 @@ bzl_library( "@bazel_skylib//lib:versions", ], ) + +bzl_library( + name = "resource_sets", + srcs = ["resource_sets.bzl"], + visibility = ["//oci:__subpackages__"], +) diff --git a/oci/private/image.bzl b/oci/private/image.bzl index f73e6edc..55a377b2 100644 --- a/oci/private/image.bzl +++ b/oci/private/image.bzl @@ -1,6 +1,7 @@ "Implementation details for image rule" -load("//oci/private:util.bzl", "util") +load("resource_sets.bzl", "resource_set", "resource_set_attr") +load("util.bzl", "util") _DOC = """Build an OCI compatible container image. @@ -189,6 +190,7 @@ def _oci_image_impl(ctx): tools = [crane.crane_info.binary, registry.registry_info.launcher, registry.registry_info.registry, jq.jqinfo.bin], mnemonic = "OCIImage", progress_message = "OCI Image %{label}", + resource_set = resource_set(ctx.attr), ) return [ @@ -199,7 +201,7 @@ def _oci_image_impl(ctx): oci_image = rule( implementation = _oci_image_impl, - attrs = _attrs, + attrs = dict(_attrs, **resource_set_attr), doc = _DOC, toolchains = [ "@bazel_tools//tools/sh:toolchain_type", diff --git a/oci/private/resource_sets.bzl b/oci/private/resource_sets.bzl new file mode 100644 index 00000000..ac1ce7fe --- /dev/null +++ b/oci/private/resource_sets.bzl @@ -0,0 +1,89 @@ +"""Utilities for rules that expose resource_set on ctx.actions.run[_shell] + +Workaround for https://github.com/bazelbuild/bazel/issues/15187 +Vendored from https://github.com/aspect-build/bazel-lib/blob/cc956d8589c866339f81637037435366c25f582b/lib/resource_sets.bzl + +Note, this workaround only provides some fixed values for either CPU or Memory. + +Rule authors who are ALSO the BUILD author might know better, and can +write custom resource_set functions for use within their own repository. +This seems to be the use case that Google engineers imagined. +""" + +resource_set_values = [ + "cpu_2", + "cpu_4", + "default", + "mem_512m", + "mem_1g", + "mem_2g", + "mem_4g", + "mem_8g", + "mem_16g", + "mem_32g", +] + +def _resource_set_cpu_2(_, __): + return {"cpu": 2} + +def _resource_set_cpu_4(_, __): + return {"cpu": 4} + +def _resource_set_mem_512m(_, __): + return {"memory": 512} + +def _resource_set_mem_1g(_, __): + return {"memory": 1024} + +def _resource_set_mem_2g(_, __): + return {"memory": 2048} + +def _resource_set_mem_4g(_, __): + return {"memory": 4096} + +def _resource_set_mem_8g(_, __): + return {"memory": 8192} + +def _resource_set_mem_16g(_, __): + return {"memory": 16384} + +def _resource_set_mem_32g(_, __): + return {"memory": 32768} + +# buildifier: disable=function-docstring +def resource_set(attr): + if attr.resource_set == "cpu_2": + return _resource_set_cpu_2 + if attr.resource_set == "cpu_4": + return _resource_set_cpu_4 + if attr.resource_set == "default": + return None + if attr.resource_set == "mem_512m": + return _resource_set_mem_512m + if attr.resource_set == "mem_1g": + return _resource_set_mem_1g + if attr.resource_set == "mem_2g": + return _resource_set_mem_2g + if attr.resource_set == "mem_4g": + return _resource_set_mem_4g + if attr.resource_set == "mem_8g": + return _resource_set_mem_8g + if attr.resource_set == "mem_16g": + return _resource_set_mem_16g + if attr.resource_set == "mem_32g": + return _resource_set_mem_32g + fail("unknown resource set", attr.resource_set) + +resource_set_attr = { + "resource_set": attr.string( + doc = """A predefined function used as the resource_set for actions. + + Used with --experimental_action_resource_set to reserve more RAM/CPU, preventing Bazel overscheduling resource-intensive actions. + + By default, Bazel allocates 1 CPU and 250M of RAM. + https://github.com/bazelbuild/bazel/blob/058f943037e21710837eda9ca2f85b5f8538c8c5/src/main/java/com/google/devtools/build/lib/actions/AbstractAction.java#L77 + """, + default = "default", + values = resource_set_values, + ), +}