feat: provide affordance for OCIImage actions to reserve more resourc…

…es (#650)

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)],
 )
-

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__"],
+)

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",

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,
+    ),
+}