From 828ae6fe69512ef8821985cb8778f213e4f64266 Mon Sep 17 00:00:00 2001 From: "Ian (Hee) Cha" Date: Tue, 6 Jun 2023 10:43:42 -0700 Subject: [PATCH] Add ServerCapabilities into RemoteExecutionClient (#18442) In #18202, we discussed the possibility of conditionally using the new field exclusively based on the Remote Execution Server's capabilities. Capture Remote Execution Server's capabilities and store it in RemoteExecutor implementations for furture usage. Closes #18269. PiperOrigin-RevId: 531999688 Change-Id: I370869a45c804af1ec499b9c1654c6977c7ab7d0 Co-authored-by: Son Luong Ngoc --- .../ExperimentalGrpcRemoteExecutor.java | 9 +++++++ .../build/lib/remote/GrpcRemoteExecutor.java | 9 +++++++ .../build/lib/remote/RemoteModule.java | 13 +++++++--- .../remote/common/RemoteExecutionClient.java | 4 +++ .../ExperimentalGrpcRemoteExecutorTest.java | 10 ++++++- ...SpawnRunnerWithGrpcRemoteExecutorTest.java | 26 ++++++++++++------- 6 files changed, 58 insertions(+), 13 deletions(-) diff --git a/src/main/java/com/google/devtools/build/lib/remote/ExperimentalGrpcRemoteExecutor.java b/src/main/java/com/google/devtools/build/lib/remote/ExperimentalGrpcRemoteExecutor.java index d50a77cd1c32b2..428da883fba69d 100644 --- a/src/main/java/com/google/devtools/build/lib/remote/ExperimentalGrpcRemoteExecutor.java +++ b/src/main/java/com/google/devtools/build/lib/remote/ExperimentalGrpcRemoteExecutor.java @@ -20,6 +20,7 @@ import build.bazel.remote.execution.v2.ExecutionGrpc; import build.bazel.remote.execution.v2.ExecutionGrpc.ExecutionBlockingStub; import build.bazel.remote.execution.v2.RequestMetadata; +import build.bazel.remote.execution.v2.ServerCapabilities; import build.bazel.remote.execution.v2.WaitExecutionRequest; import com.google.common.base.Preconditions; import com.google.devtools.build.lib.authandtls.CallCredentialsProvider; @@ -56,6 +57,7 @@ @ThreadSafe public class ExperimentalGrpcRemoteExecutor implements RemoteExecutionClient { + private final ServerCapabilities serverCapabilities; private final RemoteOptions remoteOptions; private final ReferenceCountedChannel channel; private final CallCredentialsProvider callCredentialsProvider; @@ -64,10 +66,12 @@ public class ExperimentalGrpcRemoteExecutor implements RemoteExecutionClient { private final AtomicBoolean closed = new AtomicBoolean(); public ExperimentalGrpcRemoteExecutor( + ServerCapabilities serverCapabilities, RemoteOptions remoteOptions, ReferenceCountedChannel channel, CallCredentialsProvider callCredentialsProvider, RemoteRetrier retrier) { + this.serverCapabilities = serverCapabilities; this.remoteOptions = remoteOptions; this.channel = channel; this.callCredentialsProvider = callCredentialsProvider; @@ -318,6 +322,11 @@ static ExecuteResponse extractResponseOrThrowIfError(Operation operation) throws } } + @Override + public ServerCapabilities getServerCapabilities() { + return this.serverCapabilities; + } + @Override public ExecuteResponse executeRemotely( RemoteActionExecutionContext context, ExecuteRequest request, OperationObserver observer) diff --git a/src/main/java/com/google/devtools/build/lib/remote/GrpcRemoteExecutor.java b/src/main/java/com/google/devtools/build/lib/remote/GrpcRemoteExecutor.java index 8cd463a84ab74d..97abfd36eefbf2 100644 --- a/src/main/java/com/google/devtools/build/lib/remote/GrpcRemoteExecutor.java +++ b/src/main/java/com/google/devtools/build/lib/remote/GrpcRemoteExecutor.java @@ -19,6 +19,7 @@ import build.bazel.remote.execution.v2.ExecutionGrpc; import build.bazel.remote.execution.v2.ExecutionGrpc.ExecutionBlockingStub; import build.bazel.remote.execution.v2.RequestMetadata; +import build.bazel.remote.execution.v2.ServerCapabilities; import build.bazel.remote.execution.v2.WaitExecutionRequest; import com.google.common.base.Preconditions; import com.google.devtools.build.lib.authandtls.CallCredentialsProvider; @@ -43,6 +44,7 @@ @ThreadSafe class GrpcRemoteExecutor implements RemoteExecutionClient { + private final ServerCapabilities serverCapabilities; private final ReferenceCountedChannel channel; private final CallCredentialsProvider callCredentialsProvider; private final RemoteRetrier retrier; @@ -50,9 +52,11 @@ class GrpcRemoteExecutor implements RemoteExecutionClient { private final AtomicBoolean closed = new AtomicBoolean(); public GrpcRemoteExecutor( + ServerCapabilities serverCapabilities, ReferenceCountedChannel channel, CallCredentialsProvider callCredentialsProvider, RemoteRetrier retrier) { + this.serverCapabilities = serverCapabilities; this.channel = channel; this.callCredentialsProvider = callCredentialsProvider; this.retrier = retrier; @@ -89,6 +93,11 @@ private ExecuteResponse getOperationResponse(Operation op) throws IOException { return null; } + @Override + public ServerCapabilities getServerCapabilities() { + return this.serverCapabilities; + } + /* Execute has two components: the Execute call and (optionally) the WaitExecution call. * This is the simple flow without any errors: * diff --git a/src/main/java/com/google/devtools/build/lib/remote/RemoteModule.java b/src/main/java/com/google/devtools/build/lib/remote/RemoteModule.java index d0afeb2d27e00e..2795b133a1224f 100644 --- a/src/main/java/com/google/devtools/build/lib/remote/RemoteModule.java +++ b/src/main/java/com/google/devtools/build/lib/remote/RemoteModule.java @@ -489,11 +489,12 @@ public void beforeCommand(CommandEnvironment env) throws AbruptExitException { // // If they point to different endpoints, we check the endpoint with execution or cache // capabilities respectively. + ServerCapabilities executionCapabilities = null; ServerCapabilities cacheCapabilities = null; try { if (execChannel != null) { if (cacheChannel != execChannel) { - var unused = + executionCapabilities = getAndVerifyServerCapabilities( remoteOptions, execChannel, @@ -521,6 +522,7 @@ public void beforeCommand(CommandEnvironment env) throws AbruptExitException { env, digestUtil, ServerCapabilitiesRequirement.EXECUTION_AND_CACHE); + executionCapabilities = cacheCapabilities; } } else { cacheCapabilities = @@ -601,7 +603,11 @@ public void beforeCommand(CommandEnvironment env) throws AbruptExitException { circuitBreaker); remoteExecutor = new ExperimentalGrpcRemoteExecutor( - remoteOptions, execChannel.retain(), callCredentialsProvider, execRetrier); + executionCapabilities, + remoteOptions, + execChannel.retain(), + callCredentialsProvider, + execRetrier); } else { RemoteRetrier execRetrier = new RemoteRetrier( @@ -610,7 +616,8 @@ public void beforeCommand(CommandEnvironment env) throws AbruptExitException { retryScheduler, circuitBreaker); remoteExecutor = - new GrpcRemoteExecutor(execChannel.retain(), callCredentialsProvider, execRetrier); + new GrpcRemoteExecutor( + executionCapabilities, execChannel.retain(), callCredentialsProvider, execRetrier); } execChannel.release(); RemoteExecutionCache remoteCache = diff --git a/src/main/java/com/google/devtools/build/lib/remote/common/RemoteExecutionClient.java b/src/main/java/com/google/devtools/build/lib/remote/common/RemoteExecutionClient.java index ac2c283ee8ff2d..eff0c581dcbbac 100644 --- a/src/main/java/com/google/devtools/build/lib/remote/common/RemoteExecutionClient.java +++ b/src/main/java/com/google/devtools/build/lib/remote/common/RemoteExecutionClient.java @@ -15,6 +15,7 @@ import build.bazel.remote.execution.v2.ExecuteRequest; import build.bazel.remote.execution.v2.ExecuteResponse; +import build.bazel.remote.execution.v2.ServerCapabilities; import java.io.IOException; /** @@ -24,6 +25,9 @@ */ public interface RemoteExecutionClient { + /** Returns the cache capabilities of the remote execution server */ + ServerCapabilities getServerCapabilities(); + /** Execute an action remotely using Remote Execution API. */ ExecuteResponse executeRemotely( RemoteActionExecutionContext context, ExecuteRequest request, OperationObserver observer) diff --git a/src/test/java/com/google/devtools/build/lib/remote/ExperimentalGrpcRemoteExecutorTest.java b/src/test/java/com/google/devtools/build/lib/remote/ExperimentalGrpcRemoteExecutorTest.java index ff3b6c79a732ef..d42021a4027141 100644 --- a/src/test/java/com/google/devtools/build/lib/remote/ExperimentalGrpcRemoteExecutorTest.java +++ b/src/test/java/com/google/devtools/build/lib/remote/ExperimentalGrpcRemoteExecutorTest.java @@ -21,8 +21,10 @@ import build.bazel.remote.execution.v2.Digest; import build.bazel.remote.execution.v2.ExecuteRequest; import build.bazel.remote.execution.v2.ExecuteResponse; +import build.bazel.remote.execution.v2.ExecutionCapabilities; import build.bazel.remote.execution.v2.OutputFile; import build.bazel.remote.execution.v2.RequestMetadata; +import build.bazel.remote.execution.v2.ServerCapabilities; import com.google.common.util.concurrent.ListeningScheduledExecutorService; import com.google.common.util.concurrent.MoreExecutors; import com.google.devtools.build.lib.authandtls.CallCredentialsProvider; @@ -133,9 +135,15 @@ public int maxConcurrency() { } }); + ServerCapabilities caps = + ServerCapabilities.newBuilder() + .setExecutionCapabilities( + ExecutionCapabilities.newBuilder().setExecEnabled(true).build()) + .build(); + executor = new ExperimentalGrpcRemoteExecutor( - remoteOptions, channel, CallCredentialsProvider.NO_CREDENTIALS, retrier); + caps, remoteOptions, channel, CallCredentialsProvider.NO_CREDENTIALS, retrier); } @After diff --git a/src/test/java/com/google/devtools/build/lib/remote/RemoteSpawnRunnerWithGrpcRemoteExecutorTest.java b/src/test/java/com/google/devtools/build/lib/remote/RemoteSpawnRunnerWithGrpcRemoteExecutorTest.java index 7088859f7653da..8c3c003826eca3 100644 --- a/src/test/java/com/google/devtools/build/lib/remote/RemoteSpawnRunnerWithGrpcRemoteExecutorTest.java +++ b/src/test/java/com/google/devtools/build/lib/remote/RemoteSpawnRunnerWithGrpcRemoteExecutorTest.java @@ -34,6 +34,7 @@ import build.bazel.remote.execution.v2.Directory; import build.bazel.remote.execution.v2.ExecuteRequest; import build.bazel.remote.execution.v2.ExecuteResponse; +import build.bazel.remote.execution.v2.ExecutionCapabilities; import build.bazel.remote.execution.v2.ExecutionGrpc.ExecutionImplBase; import build.bazel.remote.execution.v2.FileNode; import build.bazel.remote.execution.v2.FindMissingBlobsRequest; @@ -42,6 +43,7 @@ import build.bazel.remote.execution.v2.OutputDirectory; import build.bazel.remote.execution.v2.OutputFile; import build.bazel.remote.execution.v2.RequestMetadata; +import build.bazel.remote.execution.v2.ServerCapabilities; import build.bazel.remote.execution.v2.Tree; import build.bazel.remote.execution.v2.WaitExecutionRequest; import com.google.bytestream.ByteStreamGrpc.ByteStreamImplBase; @@ -198,13 +200,13 @@ public final void setUp() throws Exception { new FakeOwner("Mnemonic", "Progress Message", "//dummy:label"), ImmutableList.of("/bin/echo", "Hi!"), ImmutableMap.of("VARIABLE", "value"), - /*executionInfo=*/ ImmutableMap.of(), - /*runfilesSupplier=*/ null, - /*filesetMappings=*/ ImmutableMap.of(), - /*inputs=*/ NestedSetBuilder.create( + /* executionInfo= */ ImmutableMap.of(), + /* runfilesSupplier= */ null, + /* filesetMappings= */ ImmutableMap.of(), + /* inputs= */ NestedSetBuilder.create( Order.STABLE_ORDER, ActionInputHelper.fromPath("input")), - /*tools=*/ NestedSetBuilder.emptySet(Order.STABLE_ORDER), - /*outputs=*/ ImmutableSet.of( + /* tools= */ NestedSetBuilder.emptySet(Order.STABLE_ORDER), + /* outputs= */ ImmutableSet.of( new ActionInput() { @Override public String getExecPathString() { @@ -247,7 +249,7 @@ public PathFragment getExecPath() { return PathFragment.create("bar"); } }), - /*mandatoryOutputs=*/ ImmutableSet.of(), + /* mandatoryOutputs= */ ImmutableSet.of(), ResourceSet.ZERO); Path stdout = fs.getPath("/tmp/stdout"); @@ -295,8 +297,14 @@ public int maxConcurrency() { return 100; } }); + ServerCapabilities caps = + ServerCapabilities.newBuilder() + .setExecutionCapabilities( + ExecutionCapabilities.newBuilder().setExecEnabled(true).build()) + .build(); GrpcRemoteExecutor executor = - new GrpcRemoteExecutor(channel.retain(), CallCredentialsProvider.NO_CREDENTIALS, retrier); + new GrpcRemoteExecutor( + caps, channel.retain(), CallCredentialsProvider.NO_CREDENTIALS, retrier); CallCredentialsProvider callCredentialsProvider = GoogleAuthUtils.newCallCredentialsProvider(null); GrpcCacheClient cacheProtocol = @@ -326,7 +334,7 @@ public int maxConcurrency() { remoteOptions, Options.getDefaults(ExecutionOptions.class), /* verboseFailures= */ true, - /*cmdlineReporter=*/ null, + /* cmdlineReporter= */ null, retryService, logDir, remoteExecutionService);