89 lines
2.5 KiB
Python
89 lines
2.5 KiB
Python
|
"""Utilities for rules that expose resource_set on ctx.actions.run[_shell]
|
||
|
|
|
||
|
|
Workaround for https://github.com/bazelbuild/bazel/issues/15187
|
||
|
|
|
||
|
|
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,
|
||
|
|
),
|
||
|
|
}
|