Skip to content

Commit

Permalink
Add ServerCapabilities into RemoteExecutionClient
Browse files Browse the repository at this point in the history
In bazelbuild#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 bazelbuild#18269.

PiperOrigin-RevId: 531999688
Change-Id: I370869a45c804af1ec499b9c1654c6977c7ab7d0
  • Loading branch information
sluongng authored and iancha1992 committed May 17, 2023
1 parent bf82e7b commit dbb9feb
Show file tree
Hide file tree
Showing 6 changed files with 58 additions and 13 deletions.
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down Expand Up @@ -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;
Expand All @@ -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;
Expand Down Expand Up @@ -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)
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand All @@ -43,16 +44,19 @@
@ThreadSafe
class GrpcRemoteExecutor implements RemoteExecutionClient {

private final ServerCapabilities serverCapabilities;
private final ReferenceCountedChannel channel;
private final CallCredentialsProvider callCredentialsProvider;
private final RemoteRetrier retrier;

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;
Expand Down Expand Up @@ -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:
*
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down Expand Up @@ -521,6 +522,7 @@ public void beforeCommand(CommandEnvironment env) throws AbruptExitException {
env,
digestUtil,
ServerCapabilitiesRequirement.EXECUTION_AND_CACHE);
executionCapabilities = cacheCapabilities;
}
} else {
cacheCapabilities =
Expand Down Expand Up @@ -601,7 +603,11 @@ public void beforeCommand(CommandEnvironment env) throws AbruptExitException {
Retrier.ALLOW_ALL_CALLS);
remoteExecutor =
new ExperimentalGrpcRemoteExecutor(
remoteOptions, execChannel.retain(), callCredentialsProvider, execRetrier);
executionCapabilities,
remoteOptions,
execChannel.retain(),
callCredentialsProvider,
execRetrier);
} else {
RemoteRetrier execRetrier =
new RemoteRetrier(
Expand All @@ -610,7 +616,8 @@ public void beforeCommand(CommandEnvironment env) throws AbruptExitException {
retryScheduler,
Retrier.ALLOW_ALL_CALLS);
remoteExecutor =
new GrpcRemoteExecutor(execChannel.retain(), callCredentialsProvider, execRetrier);
new GrpcRemoteExecutor(
executionCapabilities, execChannel.retain(), callCredentialsProvider, execRetrier);
}
execChannel.release();
RemoteExecutionCache remoteCache =
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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;

/**
Expand All @@ -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)
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down Expand Up @@ -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
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand All @@ -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;
Expand Down Expand Up @@ -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.<String, String>of(),
/*runfilesSupplier=*/ null,
/*filesetMappings=*/ ImmutableMap.of(),
/*inputs=*/ NestedSetBuilder.create(
/* executionInfo= */ ImmutableMap.<String, String>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() {
Expand Down Expand Up @@ -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");
Expand Down Expand Up @@ -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 =
Expand Down Expand Up @@ -326,7 +334,7 @@ public int maxConcurrency() {
remoteOptions,
Options.getDefaults(ExecutionOptions.class),
/* verboseFailures= */ true,
/*cmdlineReporter=*/ null,
/* cmdlineReporter= */ null,
retryService,
logDir,
remoteExecutionService);
Expand Down

0 comments on commit dbb9feb

Please sign in to comment.