Include worker-key-mnemonic when copying the execution_requirements into the execution info.
This key allows the sharing of workers between actions with different mnemonics. Note the problematic filtering occurs in registerStarlarkAction in StarlarkActionFactory.java.
PiperOrigin-RevId: 317141655
diff --git a/src/main/java/com/google/devtools/build/lib/packages/BUILD b/src/main/java/com/google/devtools/build/lib/packages/BUILD
index 93313ad..03d8627 100644
--- a/src/main/java/com/google/devtools/build/lib/packages/BUILD
+++ b/src/main/java/com/google/devtools/build/lib/packages/BUILD
@@ -33,6 +33,7 @@
deps = [
":build_type",
":type",
+ "//src/main/java/com/google/devtools/build/lib/actions:execution_requirements",
"//src/main/java/com/google/devtools/build/lib/analysis:config/config_matching_provider",
"//src/main/java/com/google/devtools/build/lib/analysis:config/transitions/configuration_transition",
"//src/main/java/com/google/devtools/build/lib/analysis:config/transitions/no_transition",
@@ -65,7 +66,6 @@
"//src/main/java/net/starlark/java/spelling",
"//src/main/protobuf:build_java_proto",
"//third_party:auto_value",
- "//third_party:flogger",
"//third_party:guava",
"//third_party:jsr305",
"//third_party/protobuf:protobuf_java",
diff --git a/src/main/java/com/google/devtools/build/lib/packages/TargetUtils.java b/src/main/java/com/google/devtools/build/lib/packages/TargetUtils.java
index 2c964ef..25e053a 100644
--- a/src/main/java/com/google/devtools/build/lib/packages/TargetUtils.java
+++ b/src/main/java/com/google/devtools/build/lib/packages/TargetUtils.java
@@ -22,6 +22,7 @@
import com.google.common.base.Strings;
import com.google.common.collect.ImmutableMap;
import com.google.common.collect.Maps;
+import com.google.devtools.build.lib.actions.ExecutionRequirements;
import com.google.devtools.build.lib.cmdline.Label;
import com.google.devtools.build.lib.syntax.Dict;
import com.google.devtools.build.lib.syntax.EvalException;
@@ -53,8 +54,9 @@
|| tag.startsWith("no-")
|| tag.startsWith("supports-")
|| tag.startsWith("disable-")
- || tag.equals("local")
- || tag.startsWith("cpu:");
+ || tag.startsWith("cpu:")
+ || tag.equals(ExecutionRequirements.LOCAL)
+ || tag.equals(ExecutionRequirements.WORKER_KEY_MNEMONIC);
}
private TargetUtils() {} // Uninstantiable.
diff --git a/src/test/java/com/google/devtools/build/lib/packages/TargetUtilsTest.java b/src/test/java/com/google/devtools/build/lib/packages/TargetUtilsTest.java
index d140af0..bfc69cf 100644
--- a/src/test/java/com/google/devtools/build/lib/packages/TargetUtilsTest.java
+++ b/src/test/java/com/google/devtools/build/lib/packages/TargetUtilsTest.java
@@ -191,7 +191,7 @@
}
@Test
- public void testFilteredExecutionInfo_FromUncheckedExecRequirements() throws Exception {
+ public void testFilteredExecutionInfo_fromUncheckedExecRequirements() throws Exception {
scratch.file("tests/BUILD", "sh_binary(name = 'no-tag', srcs=['sh.sh'])");
Rule noTag = (Rule) getTarget("//tests:no-tag");
@@ -212,6 +212,23 @@
}
@Test
+ public void testFilteredExecutionInfo_fromUncheckedExecRequirements_withWorkerKeyMnemonic()
+ throws Exception {
+ scratch.file("tests/BUILD", "sh_binary(name = 'no-tag', srcs=['sh.sh'])");
+
+ Rule noTag = (Rule) getTarget("//tests:no-tag");
+
+ Map<String, String> execInfo =
+ TargetUtils.getFilteredExecutionInfo(
+ Dict.of(
+ (Mutability) null, "supports-workers", "1", "worker-key-mnemonic", "MyMnemonic"),
+ noTag, /* allowTagsPropagation */
+ true);
+ assertThat(execInfo)
+ .containsExactly("supports-workers", "1", "worker-key-mnemonic", "MyMnemonic");
+ }
+
+ @Test
public void testFilteredExecutionInfo() throws Exception {
scratch.file(
"tests/BUILD",
@@ -244,7 +261,7 @@
}
@Test
- public void testFilteredExecutionInfo_WithNullUncheckedExecRequirements() throws Exception {
+ public void testFilteredExecutionInfo_withNullUncheckedExecRequirements() throws Exception {
scratch.file(
"tests/BUILD",
"sh_binary(name = 'tag1', srcs=['sh.sh'], tags=['supports-workers', 'no-cache'])");
@@ -260,7 +277,7 @@
}
@Test
- public void testFilteredExecutionInfoWhenIncompatibleFlagDisabled() throws Exception {
+ public void testFilteredExecutionInfo_whenIncompatibleFlagDisabled() throws Exception {
// when --incompatible_allow_tags_propagation=false
scratch.file(
"tests/BUILD",
diff --git a/src/test/java/com/google/devtools/build/lib/skylark/BUILD b/src/test/java/com/google/devtools/build/lib/skylark/BUILD
index 30ec4a2..531ed07 100644
--- a/src/test/java/com/google/devtools/build/lib/skylark/BUILD
+++ b/src/test/java/com/google/devtools/build/lib/skylark/BUILD
@@ -58,15 +58,12 @@
"//src/main/java/com/google/devtools/build/lib/rules/python",
"//src/main/java/com/google/devtools/build/lib/skyframe:aspect_value_key",
"//src/main/java/com/google/devtools/build/lib/skyframe:configured_target_and_data",
- "//src/main/java/com/google/devtools/build/lib/skyframe:sky_functions",
"//src/main/java/com/google/devtools/build/lib/skyframe:skyframe_cluster",
"//src/main/java/com/google/devtools/build/lib/util",
"//src/main/java/com/google/devtools/build/lib/util:classpath",
"//src/main/java/com/google/devtools/build/lib/util:filetype",
"//src/main/java/com/google/devtools/build/lib/vfs",
"//src/main/java/com/google/devtools/build/lib/vfs:pathfragment",
- "//src/main/java/com/google/devtools/build/skyframe",
- "//src/main/java/com/google/devtools/build/skyframe:skyframe-objects",
"//src/main/java/com/google/devtools/common/options",
"//src/main/java/net/starlark/java/annot",
"//src/test/java/com/google/devtools/build/lib/actions/util",
diff --git a/src/test/java/com/google/devtools/build/lib/skylark/StarlarkRuleImplementationFunctionsTest.java b/src/test/java/com/google/devtools/build/lib/skylark/StarlarkRuleImplementationFunctionsTest.java
index 2da9692..656d03d 100644
--- a/src/test/java/com/google/devtools/build/lib/skylark/StarlarkRuleImplementationFunctionsTest.java
+++ b/src/test/java/com/google/devtools/build/lib/skylark/StarlarkRuleImplementationFunctionsTest.java
@@ -384,6 +384,30 @@
}
@Test
+ public void testCreateSpawnActionEnvAndExecInfo_withWorkerKeyMnemonic() throws Exception {
+ StarlarkRuleContext ruleContext = createRuleContext("//foo:foo");
+ setRuleContext(ruleContext);
+ ev.exec(
+ "ruleContext.actions.run_shell(",
+ " inputs = ruleContext.files.srcs,",
+ " outputs = ruleContext.files.srcs,",
+ " env = {'a' : 'b'},",
+ " execution_requirements = {",
+ " 'supports-workers': '1',",
+ " 'worker-key-mnemonic': 'MyMnemonic',",
+ " },",
+ " mnemonic = 'DummyMnemonic',",
+ " command = 'dummy_command',",
+ " progress_message = 'dummy_message')");
+ SpawnAction action =
+ (SpawnAction)
+ Iterables.getOnlyElement(
+ ruleContext.getRuleContext().getAnalysisEnvironment().getRegisteredActions());
+ assertThat(action.getExecutionInfo())
+ .containsExactly("supports-workers", "1", "worker-key-mnemonic", "MyMnemonic");
+ }
+
+ @Test
public void testCreateSpawnActionUnknownParam() throws Exception {
setRuleContext(createRuleContext("//foo:foo"));
ev.checkEvalErrorContains(
diff --git a/src/test/shell/integration/bazel_worker_test.sh b/src/test/shell/integration/bazel_worker_test.sh
index 0714bf1..2d51904 100755
--- a/src/test/shell/integration/bazel_worker_test.sh
+++ b/src/test/shell/integration/bazel_worker_test.sh
@@ -126,13 +126,17 @@
argfile_inputs.append(argfile)
argfile_arguments.append("@" + argfile.path)
+ execution_requirements = {"supports-workers": "1"}
+ if ctx.attr.worker_key_mnemonic:
+ execution_requirements["worker-key-mnemonic"] = ctx.attr.worker_key_mnemonic
+
ctx.actions.run(
inputs=argfile_inputs + ctx.files.srcs,
outputs=[output],
executable=worker,
progress_message="Working on %s" % ctx.label.name,
- mnemonic="Work",
- execution_requirements={"supports-workers": "1"},
+ mnemonic=ctx.attr.action_mnemonic,
+ execution_requirements=execution_requirements,
arguments=ctx.attr.worker_args + argfile_arguments,
)
@@ -141,6 +145,8 @@
attrs={
"worker": attr.label(cfg="host", mandatory=True, allow_files=True, executable=True),
"worker_args": attr.string_list(),
+ "worker_key_mnemonic": attr.string(),
+ "action_mnemonic": attr.string(default = "Work"),
"args": attr.string_list(),
"srcs": attr.label_list(allow_files=True),
"multiflagfiles": attr.bool(default=False),
@@ -195,6 +201,33 @@
assert_equals "HELLO WORLD" "$(cat $BINS/hello_world_uppercase.out)"
}
+function test_shared_worker() {
+ prepare_example_worker
+ cat >>BUILD <<EOF
+work(
+ name = "hello_world",
+ worker = ":worker",
+ action_mnemonic = "Hello",
+ worker_key_mnemonic = "SharedWorker",
+ args = ["--write_uuid"],
+)
+
+work(
+ name = "goodbye_world",
+ worker = ":worker",
+ action_mnemonic = "Goodbye",
+ worker_key_mnemonic = "SharedWorker",
+ args = ["--write_uuid"],
+)
+EOF
+
+ bazel build :hello_world :goodbye_world &> $TEST_log \
+ || fail "build failed"
+ worker_uuid_1=$(cat $BINS/hello_world.out | grep UUID | cut -d' ' -f2)
+ worker_uuid_2=$(cat $BINS/goodbye_world.out | grep UUID | cut -d' ' -f2)
+ assert_equals "$worker_uuid_1" "$worker_uuid_2"
+}
+
function test_multiple_flagfiles() {
prepare_example_worker
cat >>BUILD <<EOF
@@ -456,7 +489,7 @@
)
EOF
- sed -i.bak '/execution_requirements/d' work.bzl
+ sed -i.bak '/execution_requirements=execution_requirements/d' work.bzl
rm -f work.bzl.bak
bazel build --worker_quit_after_build :hello_world &> $TEST_log \
