Skip to content

Commit d7f0724

Browse files
wilwellcopybara-github
wilwell
authored and
copybara-github
committed
ResourceSet in StarlarkAction API
Added optional `resource_set` parameter to `run` and `run_shell` in StarlarkActionApi. `resource_set` is `StarlarkCallable` object that returns dict with resource set (cpu, memory, local_test). PiperOrigin-RevId: 415224490
1 parent 44c30ac commit d7f0724

File tree

17 files changed

+465
-31
lines changed

17 files changed

+465
-31
lines changed

‎src/main/java/com/google/devtools/build/lib/actions/BUILD‎

Lines changed: 5 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -300,8 +300,7 @@ java_library(
300300
"ResourceSetOrBuilder.java",
301301
],
302302
deps = [
303-
":artifacts",
304-
"//src/main/java/com/google/devtools/build/lib/collect/nestedset",
303+
":exec_exception",
305304
"//src/main/java/com/google/devtools/build/lib/concurrent",
306305
"//src/main/java/com/google/devtools/build/lib/jni",
307306
"//src/main/java/com/google/devtools/build/lib/unix",
@@ -349,6 +348,9 @@ java_library(
349348

350349
java_library(
351350
name = "exec_exception",
352-
srcs = ["ExecException.java"],
351+
srcs = [
352+
"ExecException.java",
353+
"UserExecException.java",
354+
],
353355
deps = ["//src/main/protobuf:failure_details_java_proto"],
354356
)

‎src/main/java/com/google/devtools/build/lib/actions/BaseSpawn.java‎

Lines changed: 6 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -19,6 +19,7 @@
1919
import com.google.common.collect.Iterables;
2020
import com.google.devtools.build.lib.analysis.platform.PlatformInfo;
2121
import com.google.devtools.build.lib.collect.nestedset.NestedSet;
22+
import com.google.devtools.build.lib.util.OS;
2223
import com.google.devtools.build.lib.vfs.PathFragment;
2324
import java.util.Collection;
2425
import java.util.List;
@@ -35,15 +36,15 @@ public class BaseSpawn implements Spawn {
3536
private final ImmutableMap<String, String> executionInfo;
3637
private final RunfilesSupplier runfilesSupplier;
3738
private final ActionExecutionMetadata action;
38-
private final ResourceSet localResources;
39+
private final ResourceSetOrBuilder localResources;
3940

4041
public BaseSpawn(
4142
List<String> arguments,
4243
Map<String, String> environment,
4344
Map<String, String> executionInfo,
4445
RunfilesSupplier runfilesSupplier,
4546
ActionExecutionMetadata action,
46-
ResourceSet localResources) {
47+
ResourceSetOrBuilder localResources) {
4748
this.arguments = ImmutableList.copyOf(arguments);
4849
this.environment = ImmutableMap.copyOf(environment);
4950
this.executionInfo = ImmutableMap.copyOf(executionInfo);
@@ -123,8 +124,9 @@ public ActionExecutionMetadata getResourceOwner() {
123124
}
124125

125126
@Override
126-
public ResourceSet getLocalResources() {
127-
return localResources;
127+
public ResourceSet getLocalResources() throws ExecException {
128+
return localResources.buildResourceSet(
129+
OS.getCurrent(), action.getInputs().memoizedFlattenAndGetSize());
128130
}
129131

130132
@Override

‎src/main/java/com/google/devtools/build/lib/actions/DelegateSpawn.java‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -79,7 +79,7 @@ public ActionExecutionMetadata getResourceOwner() {
7979
}
8080

8181
@Override
82-
public ResourceSet getLocalResources() {
82+
public ResourceSet getLocalResources() throws ExecException {
8383
return spawn.getLocalResources();
8484
}
8585

‎src/main/java/com/google/devtools/build/lib/actions/ResourceSet.java‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -16,8 +16,8 @@
1616

1717
import com.google.common.base.Splitter;
1818
import com.google.common.primitives.Doubles;
19-
import com.google.devtools.build.lib.collect.nestedset.NestedSet;
2019
import com.google.devtools.build.lib.concurrent.ThreadSafety.Immutable;
20+
import com.google.devtools.build.lib.util.OS;
2121
import com.google.devtools.common.options.Converter;
2222
import com.google.devtools.common.options.OptionsParsingException;
2323
import java.util.Iterator;
@@ -175,7 +175,7 @@ public String getTypeDescription() {
175175
}
176176

177177
@Override
178-
public ResourceSet buildResourceSet(NestedSet<Artifact> inputs) {
178+
public ResourceSet buildResourceSet(OS os, int inputsSize) throws ExecException {
179179
return this;
180180
}
181181
}

‎src/main/java/com/google/devtools/build/lib/actions/ResourceSetOrBuilder.java‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -14,7 +14,7 @@
1414

1515
package com.google.devtools.build.lib.actions;
1616

17-
import com.google.devtools.build.lib.collect.nestedset.NestedSet;
17+
import com.google.devtools.build.lib.util.OS;
1818

1919
/** Common interface for ResourceSet and builder. */
2020
@FunctionalInterface
@@ -24,5 +24,5 @@ public interface ResourceSetOrBuilder {
2424
* will flatten NestedSet. This action could create a lot of garbagge, so use it as close as
2525
* possible to the execution phase,
2626
*/
27-
public ResourceSet buildResourceSet(NestedSet<Artifact> inputs);
27+
public ResourceSet buildResourceSet(OS os, int inputsSize) throws ExecException;
2828
}

‎src/main/java/com/google/devtools/build/lib/actions/Spawn.java‎

Lines changed: 2 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -115,10 +115,8 @@ public interface Spawn extends DescribableExecutionUnit {
115115
*/
116116
ActionExecutionMetadata getResourceOwner();
117117

118-
/**
119-
* Returns the amount of resources needed for local fallback.
120-
*/
121-
ResourceSet getLocalResources();
118+
/** Returns the amount of resources needed for local fallback. */
119+
ResourceSet getLocalResources() throws ExecException;
122120

123121
/**
124122
* Returns a mnemonic (string constant) for this kind of spawn.

‎src/main/java/com/google/devtools/build/lib/analysis/actions/SpawnAction.java‎

Lines changed: 6 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -353,6 +353,11 @@ public final Spawn getSpawn() throws CommandLineExpansionException, InterruptedE
353353
return getSpawn(getInputs());
354354
}
355355

356+
@VisibleForTesting
357+
public ResourceSetOrBuilder getResourceSetOrBuilder() {
358+
return resourceSetOrBuilder;
359+
}
360+
356361
final Spawn getSpawn(NestedSet<Artifact> inputs)
357362
throws CommandLineExpansionException, InterruptedException {
358363
return new ActionSpawn(
@@ -579,8 +584,7 @@ private ActionSpawn(
579584
executionInfo,
580585
SpawnAction.this.getRunfilesSupplier(),
581586
SpawnAction.this,
582-
SpawnAction.this.resourceSetOrBuilder.buildResourceSet(inputs));
583-
587+
SpawnAction.this.resourceSetOrBuilder);
584588
NestedSetBuilder<ActionInput> inputsBuilder = NestedSetBuilder.stableOrder();
585589
ImmutableList<Artifact> manifests = getRunfilesSupplier().getManifests();
586590
for (Artifact input : inputs.toList()) {

‎src/main/java/com/google/devtools/build/lib/analysis/starlark/StarlarkActionFactory.java‎

Lines changed: 146 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -15,17 +15,23 @@
1515

1616
import static com.google.devtools.build.lib.packages.semantics.BuildLanguageOptions.EXPERIMENTAL_SIBLING_REPOSITORY_LAYOUT;
1717

18+
import com.google.common.base.Joiner;
1819
import com.google.common.collect.ImmutableList;
1920
import com.google.common.collect.ImmutableMap;
21+
import com.google.common.collect.Sets;
2022
import com.google.devtools.build.lib.actions.Action;
2123
import com.google.devtools.build.lib.actions.ActionAnalysisMetadata;
2224
import com.google.devtools.build.lib.actions.ActionLookupKey;
2325
import com.google.devtools.build.lib.actions.ActionRegistry;
2426
import com.google.devtools.build.lib.actions.Artifact;
2527
import com.google.devtools.build.lib.actions.ArtifactRoot;
2628
import com.google.devtools.build.lib.actions.CommandLine;
29+
import com.google.devtools.build.lib.actions.ExecException;
2730
import com.google.devtools.build.lib.actions.ParamFileInfo;
31+
import com.google.devtools.build.lib.actions.ResourceSet;
32+
import com.google.devtools.build.lib.actions.ResourceSetOrBuilder;
2833
import com.google.devtools.build.lib.actions.RunfilesSupplier;
34+
import com.google.devtools.build.lib.actions.UserExecException;
2935
import com.google.devtools.build.lib.actions.extra.ExtraActionInfo;
3036
import com.google.devtools.build.lib.actions.extra.SpawnInfo;
3137
import com.google.devtools.build.lib.analysis.BashCommandConstructor;
@@ -49,23 +55,35 @@
4955
import com.google.devtools.build.lib.collect.nestedset.Order;
5056
import com.google.devtools.build.lib.packages.TargetUtils;
5157
import com.google.devtools.build.lib.packages.semantics.BuildLanguageOptions;
58+
import com.google.devtools.build.lib.server.FailureDetails;
59+
import com.google.devtools.build.lib.server.FailureDetails.FailureDetail;
60+
import com.google.devtools.build.lib.server.FailureDetails.Interrupted;
5261
import com.google.devtools.build.lib.skyframe.serialization.autocodec.AutoCodec;
5362
import com.google.devtools.build.lib.skyframe.serialization.autocodec.SerializationConstant;
5463
import com.google.devtools.build.lib.starlarkbuildapi.FileApi;
5564
import com.google.devtools.build.lib.starlarkbuildapi.StarlarkActionFactoryApi;
65+
import com.google.devtools.build.lib.util.OS;
5666
import com.google.devtools.build.lib.vfs.PathFragment;
5767
import com.google.protobuf.GeneratedMessage;
5868
import java.nio.charset.StandardCharsets;
5969
import java.util.ArrayList;
70+
import java.util.Arrays;
71+
import java.util.HashSet;
6072
import java.util.List;
6173
import java.util.Map;
6274
import java.util.Optional;
75+
import java.util.Set;
6376
import java.util.UUID;
6477
import net.starlark.java.eval.Dict;
6578
import net.starlark.java.eval.EvalException;
79+
import net.starlark.java.eval.Mutability;
6680
import net.starlark.java.eval.Printer;
6781
import net.starlark.java.eval.Sequence;
6882
import net.starlark.java.eval.Starlark;
83+
import net.starlark.java.eval.StarlarkCallable;
84+
import net.starlark.java.eval.StarlarkFloat;
85+
import net.starlark.java.eval.StarlarkFunction;
86+
import net.starlark.java.eval.StarlarkInt;
6987
import net.starlark.java.eval.StarlarkSemantics;
7088
import net.starlark.java.eval.StarlarkThread;
7189

@@ -75,6 +93,10 @@ public class StarlarkActionFactory implements StarlarkActionFactoryApi {
7593
/** Counter for actions.run_shell helper scripts. Every script must have a unique name. */
7694
private int runShellOutputCounter = 0;
7795

96+
private static final ResourceSet DEFAULT_RESOURCE_SET = ResourceSet.createWithRamCpu(250, 1);
97+
private static final Set<String> validResources =
98+
new HashSet<>(Arrays.asList("cpu", "memory", "local_test"));
99+
78100
public StarlarkActionFactory(StarlarkRuleContext context) {
79101
this.context = context;
80102
}
@@ -339,7 +361,8 @@ public void run(
339361
Object executionRequirementsUnchecked,
340362
Object inputManifestsUnchecked,
341363
Object execGroupUnchecked,
342-
Object shadowedActionUnchecked)
364+
Object shadowedActionUnchecked,
365+
Object resourceSetUnchecked)
343366
throws EvalException {
344367
context.checkMutable("actions.run");
345368

@@ -377,6 +400,7 @@ public void run(
377400
inputManifestsUnchecked,
378401
execGroupUnchecked,
379402
shadowedActionUnchecked,
403+
resourceSetUnchecked,
380404
builder);
381405
}
382406

@@ -430,7 +454,8 @@ public void runShell(
430454
Object executionRequirementsUnchecked,
431455
Object inputManifestsUnchecked,
432456
Object execGroupUnchecked,
433-
Object shadowedActionUnchecked)
457+
Object shadowedActionUnchecked,
458+
Object resourceSetUnchecked)
434459
throws EvalException {
435460
context.checkMutable("actions.run_shell");
436461
RuleContext ruleContext = getRuleContext();
@@ -497,6 +522,7 @@ public void runShell(
497522
inputManifestsUnchecked,
498523
execGroupUnchecked,
499524
shadowedActionUnchecked,
525+
resourceSetUnchecked,
500526
builder);
501527
}
502528

@@ -543,6 +569,7 @@ private void registerStarlarkAction(
543569
Object inputManifestsUnchecked,
544570
Object execGroupUnchecked,
545571
Object shadowedActionUnchecked,
572+
Object resourceSetUnchecked,
546573
StarlarkAction.Builder builder)
547574
throws EvalException {
548575
if (inputs instanceof Sequence) {
@@ -648,10 +675,127 @@ private void registerStarlarkAction(
648675
builder.setShadowedAction(Optional.of((Action) shadowedActionUnchecked));
649676
}
650677

678+
if (getSemantics().getBool(BuildLanguageOptions.EXPERIMENTAL_ACTION_RESOURCE_SET)
679+
&& resourceSetUnchecked != Starlark.NONE) {
680+
validateResourceSetBuilder(resourceSetUnchecked);
681+
builder.setResources(
682+
new StarlarkActionResourceSetBuilder(
683+
(StarlarkFunction) resourceSetUnchecked, mnemonic, getSemantics()));
684+
}
685+
651686
// Always register the action
652687
registerAction(builder.build(ruleContext));
653688
}
654689

690+
private static class StarlarkActionResourceSetBuilder implements ResourceSetOrBuilder {
691+
private final StarlarkCallable fn;
692+
private final String mnemonic;
693+
private final StarlarkSemantics semantics;
694+
695+
public StarlarkActionResourceSetBuilder(
696+
StarlarkCallable fn, String mnemonic, StarlarkSemantics semantics) {
697+
this.fn = fn;
698+
this.mnemonic = mnemonic;
699+
this.semantics = semantics;
700+
}
701+
702+
@Override
703+
public ResourceSet buildResourceSet(OS os, int inputsSize) throws ExecException {
704+
try (Mutability mu = Mutability.create("resource_set_builder_function")) {
705+
StarlarkThread thread = new StarlarkThread(mu, semantics);
706+
StarlarkInt inputInt = StarlarkInt.of(inputsSize);
707+
Object response =
708+
Starlark.call(
709+
thread,
710+
this.fn,
711+
ImmutableList.of(os.getCanonicalName(), inputInt),
712+
ImmutableMap.of());
713+
Map<String, Object> resourceSetMapRaw =
714+
Dict.cast(response, String.class, Object.class, "resource_set");
715+
716+
if (!validResources.containsAll(resourceSetMapRaw.keySet())) {
717+
String message =
718+
String.format(
719+
"Illegal resource keys: (%s)",
720+
Joiner.on(",").join(Sets.difference(resourceSetMapRaw.keySet(), validResources)));
721+
throw new EvalException(message);
722+
}
723+
724+
return ResourceSet.create(
725+
getNumericOrDefault(resourceSetMapRaw, "memory", DEFAULT_RESOURCE_SET.getMemoryMb()),
726+
getNumericOrDefault(resourceSetMapRaw, "cpu", DEFAULT_RESOURCE_SET.getCpuUsage()),
727+
(int)
728+
getNumericOrDefault(
729+
resourceSetMapRaw,
730+
"local_test",
731+
(double) DEFAULT_RESOURCE_SET.getLocalTestCount()));
732+
} catch (EvalException e) {
733+
throw new UserExecException(
734+
FailureDetail.newBuilder()
735+
.setMessage(
736+
String.format("Could not build resources for %s. %s", mnemonic, e.getMessage()))
737+
.setStarlarkAction(
738+
FailureDetails.StarlarkAction.newBuilder()
739+
.setCode(FailureDetails.StarlarkAction.Code.STARLARK_ACTION_UNKNOWN)
740+
.build())
741+
.build());
742+
} catch (InterruptedException e) {
743+
throw new UserExecException(
744+
FailureDetail.newBuilder()
745+
.setMessage(e.getMessage())
746+
.setInterrupted(
747+
Interrupted.newBuilder().setCode(Interrupted.Code.INTERRUPTED).build())
748+
.build());
749+
}
750+
}
751+
752+
private static double getNumericOrDefault(
753+
Map<String, Object> resourceSetMap, String key, double defaultValue) throws EvalException {
754+
if (!resourceSetMap.containsKey(key)) {
755+
return defaultValue;
756+
}
757+
758+
Object value = resourceSetMap.get(key);
759+
if (value instanceof StarlarkInt) {
760+
return ((StarlarkInt) value).toDouble();
761+
}
762+
763+
if (value instanceof StarlarkFloat) {
764+
return ((StarlarkFloat) value).toDouble();
765+
}
766+
throw new EvalException(
767+
String.format(
768+
"Illegal resource value type for key %s: got %s, want int or float",
769+
key, Starlark.type(value)));
770+
}
771+
}
772+
773+
private static StarlarkFunction validateResourceSetBuilder(Object fn) throws EvalException {
774+
if (!(fn instanceof StarlarkFunction)) {
775+
throw Starlark.errorf(
776+
"resource_set should be a Starlark-defined function, but got %s instead",
777+
Starlark.type(fn));
778+
}
779+
780+
StarlarkFunction sfn = (StarlarkFunction) fn;
781+
782+
// Reject non-global functions, because arbitrary closures may cause large
783+
// analysis-phase data structures to remain live into the execution phase.
784+
// We require that the function is "global" as opposed to "not a closure"
785+
// because a global function may be closure if it refers to load bindings.
786+
// This unfortunately disallows such trivially safe non-global
787+
// functions as "lambda x: x".
788+
// See https://github.com/bazelbuild/bazel/issues/12701.
789+
if (sfn.getModule().getGlobal(sfn.getName()) != sfn) {
790+
throw Starlark.errorf(
791+
"to avoid unintended retention of analysis data structures, "
792+
+ "the resource_set function (declared at %s) must be declared "
793+
+ "by a top-level def statement",
794+
sfn.getLocation());
795+
}
796+
return (StarlarkFunction) fn;
797+
}
798+
655799
private String getMnemonic(Object mnemonicUnchecked) {
656800
String mnemonic = mnemonicUnchecked == Starlark.NONE ? "Action" : (String) mnemonicUnchecked;
657801
if (getRuleContext().getConfiguration().getReservedActionMnemonics().contains(mnemonic)) {

0 commit comments

Comments
 (0)