Skip to content

Commit

Permalink
Add ActionExecutionMetadata as a parameter to `ActionInputPrefetche…
Browse files Browse the repository at this point in the history
…r#prefetchFiles`.

Due to recent overhual of `prefetchFiles`, the call sites now always have the reference to `ActionExecutionMetadata`, which makes it possible to pass it to `prefetchFiles`.

The reason why we want to access `ActionExecutionMetadata` inside `prefetchFiles` is to associate downloads with actions. For example:

1. When downloading from remote cache, we want to provide action metadata to the RPCs so that remote server can associate them with actions. `actionId` was removed from remote file metadata by f62a8b9 because it is not the right action to associate with for downloads inside this context. We left the action metadata empty for RPCs until this change. #17724 (comment)
2. When building with the bytes, we are able to report download progress to UI. However, `prefetchFiles` is the major way to download when building without the bytes and we can't report download progress on action basis. #17604 added the download report for background downloads, but now, we no longer download toplevel artifacts in background (a5dde12). With this change, we are able to report download progress for each action, either when they download outputs, or prefetch inputs for building without the bytes.

PiperOrigin-RevId: 539635790
Change-Id: Ic9cc1024f5d3560ac46bc462bd549dd81712d92f
  • Loading branch information
coeuvre authored and copybara-github committed Jun 12, 2023
1 parent 0573eee commit 3b39ab8
Show file tree
Hide file tree
Showing 11 changed files with 96 additions and 82 deletions.
Original file line number Diff line number Diff line change
Expand Up @@ -34,6 +34,7 @@ public interface MetadataSupplier {
new ActionInputPrefetcher() {
@Override
public ListenableFuture<Void> prefetchFiles(
ActionExecutionMetadata action,
Iterable<? extends ActionInput> inputs,
MetadataSupplier metadataSupplier,
Priority priority) {
Expand Down Expand Up @@ -79,7 +80,10 @@ public enum Priority {
* @return future success if prefetch is finished or {@link IOException}.
*/
ListenableFuture<Void> prefetchFiles(
Iterable<? extends ActionInput> inputs, MetadataSupplier metadataSupplier, Priority priority);
ActionExecutionMetadata action,
Iterable<? extends ActionInput> inputs,
MetadataSupplier metadataSupplier,
Priority priority);

/**
* Whether the prefetcher requires the metadata for a tree artifact to be available whenever one
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -252,6 +252,7 @@ public ListenableFuture<Void> prefetchInputs()
actionExecutionContext
.getActionInputPrefetcher()
.prefetchFiles(
spawn.getResourceOwner(),
getInputMapping(PathFragment.EMPTY_FRAGMENT, /* willAccessRepeatedly= */ true)
.values(),
getInputMetadataProvider()::getInputMetadata,
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -29,6 +29,7 @@
import com.google.common.flogger.GoogleLogger;
import com.google.common.util.concurrent.ListenableFuture;
import com.google.devtools.build.lib.actions.Action;
import com.google.devtools.build.lib.actions.ActionExecutionMetadata;
import com.google.devtools.build.lib.actions.ActionInput;
import com.google.devtools.build.lib.actions.ActionInputPrefetcher;
import com.google.devtools.build.lib.actions.Artifact;
Expand Down Expand Up @@ -263,6 +264,7 @@ private boolean shouldDownloadFile(Path path, FileArtifactValue metadata) {
* @param tempPath the temporary path which the input should be written to.
*/
protected abstract ListenableFuture<Void> doDownloadFile(
ActionExecutionMetadata action,
Reporter reporter,
Path tempPath,
PathFragment execPath,
Expand All @@ -285,6 +287,7 @@ protected void prefetchVirtualActionInput(VirtualActionInput input) throws IOExc
*/
@Override
public ListenableFuture<Void> prefetchFiles(
ActionExecutionMetadata action,
Iterable<? extends ActionInput> inputs,
MetadataSupplier metadataSupplier,
Priority priority) {
Expand All @@ -308,7 +311,8 @@ public ListenableFuture<Void> prefetchFiles(

Flowable<TransferResult> transfers =
Flowable.fromIterable(files)
.flatMapSingle(input -> prefetchFile(dirCtx, metadataSupplier, input, priority));
.flatMapSingle(
input -> prefetchFile(action, dirCtx, metadataSupplier, input, priority));

Completable prefetch =
Completable.using(
Expand All @@ -318,6 +322,7 @@ public ListenableFuture<Void> prefetchFiles(
}

private Single<TransferResult> prefetchFile(
ActionExecutionMetadata action,
DirectoryContext dirCtx,
MetadataSupplier metadataSupplier,
ActionInput input,
Expand Down Expand Up @@ -347,6 +352,7 @@ private Single<TransferResult> prefetchFile(

Completable result =
downloadFileNoCheckRx(
action,
dirCtx,
execRoot.getRelative(execPath),
treeRootExecPath != null ? execRoot.getRelative(treeRootExecPath) : null,
Expand Down Expand Up @@ -421,6 +427,7 @@ private Symlink maybeGetSymlink(
}

private Completable downloadFileNoCheckRx(
ActionExecutionMetadata action,
DirectoryContext dirCtx,
Path path,
@Nullable Path treeRoot,
Expand All @@ -445,6 +452,7 @@ private Completable downloadFileNoCheckRx(
toCompletable(
() ->
doDownloadFile(
action,
reporter,
tempPath,
finalPath.relativeTo(execRoot),
Expand Down Expand Up @@ -603,7 +611,7 @@ public void finalizeAction(Action action, OutputMetadataStore outputMetadataStor
try (var s = Profiler.instance().profile(ProfilerTask.REMOTE_DOWNLOAD, "Download outputs")) {
getFromFuture(
prefetchFiles(
outputsToDownload, outputMetadataStore::getOutputMetadata, Priority.HIGH));
action, outputsToDownload, outputMetadataStore::getOutputMetadata, Priority.HIGH));
}
}
}
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -25,6 +25,7 @@
import com.google.common.annotations.VisibleForTesting;
import com.google.common.collect.ImmutableList;
import com.google.common.collect.ImmutableMap;
import com.google.devtools.build.lib.actions.ActionExecutionMetadata;
import com.google.devtools.build.lib.actions.ActionInput;
import com.google.devtools.build.lib.actions.ActionInputHelper;
import com.google.devtools.build.lib.actions.ActionInputMap;
Expand Down Expand Up @@ -81,6 +82,7 @@ public class RemoteActionFileSystem extends DelegateFileSystem {
private final RemoteActionInputFetcher inputFetcher;
private final RemoteInMemoryFileSystem remoteOutputTree;

@Nullable private ActionExecutionMetadata action = null;
@Nullable private MetadataInjector metadataInjector = null;

RemoteActionFileSystem(
Expand Down Expand Up @@ -121,7 +123,8 @@ private boolean isRemote(PathFragment path) {
return getRemoteMetadata(path) != null;
}

public void updateContext(MetadataInjector metadataInjector) {
public void updateContext(ActionExecutionMetadata action, MetadataInjector metadataInjector) {
this.action = action;
this.metadataInjector = metadataInjector;
}

Expand Down Expand Up @@ -625,7 +628,7 @@ private void downloadFileIfRemote(PathFragment path) throws IOException {
}
getFromFuture(
inputFetcher.prefetchFiles(
ImmutableList.of(input), this::getInputMetadata, Priority.CRITICAL));
action, ImmutableList.of(input), this::getInputMetadata, Priority.CRITICAL));
} catch (InterruptedException e) {
Thread.currentThread().interrupt();
throw new IOException(String.format("Received interrupt while fetching file '%s'", path), e);
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -14,18 +14,15 @@
package com.google.devtools.build.lib.remote;

import static com.google.common.base.Preconditions.checkArgument;
import static com.google.devtools.build.lib.remote.common.ProgressStatusListener.NO_ACTION;

import build.bazel.remote.execution.v2.Digest;
import build.bazel.remote.execution.v2.RequestMetadata;
import com.google.common.base.Preconditions;
import com.google.common.util.concurrent.ListenableFuture;
import com.google.devtools.build.lib.actions.ActionExecutionMetadata;
import com.google.devtools.build.lib.actions.FileArtifactValue;
import com.google.devtools.build.lib.actions.cache.VirtualActionInput;
import com.google.devtools.build.lib.events.ExtendedEventHandler.FetchProgress;
import com.google.devtools.build.lib.events.Reporter;
import com.google.devtools.build.lib.exec.SpawnProgressEvent;
import com.google.devtools.build.lib.remote.RemoteCache.DownloadProgressReporter;
import com.google.devtools.build.lib.remote.common.RemoteActionExecutionContext;
import com.google.devtools.build.lib.remote.util.DigestUtil;
import com.google.devtools.build.lib.remote.util.TempPathGenerator;
Expand Down Expand Up @@ -79,6 +76,7 @@ protected boolean canDownloadFile(Path path, FileArtifactValue metadata) {

@Override
protected ListenableFuture<Void> doDownloadFile(
ActionExecutionMetadata action,
Reporter reporter,
Path tempPath,
PathFragment execPath,
Expand All @@ -87,47 +85,18 @@ protected ListenableFuture<Void> doDownloadFile(
throws IOException {
checkArgument(metadata.isRemote(), "Cannot download file that is not a remote file.");
RequestMetadata requestMetadata =
TracingMetadataUtils.buildMetadata(buildRequestId, commandId, "prefetcher", null);
TracingMetadataUtils.buildMetadata(buildRequestId, commandId, "prefetcher", action);
RemoteActionExecutionContext context = RemoteActionExecutionContext.create(requestMetadata);

Digest digest = DigestUtil.buildDigest(metadata.getDigest(), metadata.getSize());

DownloadProgressReporter downloadProgressReporter;
if (priority == Priority.LOW) {
// Only report download progress for toplevel outputs
downloadProgressReporter =
new DownloadProgressReporter(
/* includeFile= */ false,
progress -> reporter.post(new DownloadProgress(progress)),
execPath.toString(),
metadata.getSize());
} else {
downloadProgressReporter = new DownloadProgressReporter(NO_ACTION, "", 0);
}

return remoteCache.downloadFile(context, tempPath, digest, downloadProgressReporter);
}

public static class DownloadProgress implements FetchProgress {
private final SpawnProgressEvent progress;

public DownloadProgress(SpawnProgressEvent progress) {
this.progress = progress;
}

@Override
public String getResourceIdentifier() {
return progress.progressId();
}

@Override
public String getProgress() {
return progress.progress();
}

@Override
public boolean isFinished() {
return progress.finished();
}
return remoteCache.downloadFile(
context,
tempPath,
digest,
new RemoteCache.DownloadProgressReporter(
progress -> progress.postTo(reporter, action),
execPath.toString(),
digest.getSizeBytes()));
}
}
Original file line number Diff line number Diff line change
Expand Up @@ -21,6 +21,7 @@
import com.google.common.collect.ImmutableMap;
import com.google.common.eventbus.Subscribe;
import com.google.devtools.build.lib.actions.Action;
import com.google.devtools.build.lib.actions.ActionExecutionMetadata;
import com.google.devtools.build.lib.actions.ActionInputMap;
import com.google.devtools.build.lib.actions.Artifact;
import com.google.devtools.build.lib.actions.ArtifactPathResolver;
Expand Down Expand Up @@ -99,11 +100,12 @@ public FileSystem createActionFileSystem(

@Override
public void updateActionFileSystemContext(
ActionExecutionMetadata action,
FileSystem actionFileSystem,
Environment env,
MetadataInjector injector,
ImmutableMap<Artifact, ImmutableList<FilesetOutputSymlink>> filesets) {
((RemoteActionFileSystem) actionFileSystem).updateContext(injector);
((RemoteActionFileSystem) actionFileSystem).updateContext(action, injector);
}

@Override
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -382,11 +382,13 @@ FileSystem createActionFileSystem(
}

private void updateActionFileSystemContext(
Action action,
FileSystem actionFileSystem,
Environment env,
MetadataInjector metadataInjector,
ImmutableMap<Artifact, ImmutableList<FilesetOutputSymlink>> filesets) {
outputService.updateActionFileSystemContext(actionFileSystem, env, metadataInjector, filesets);
outputService.updateActionFileSystemContext(
action, actionFileSystem, env, metadataInjector, filesets);
}

void executionOver() {
Expand Down Expand Up @@ -489,7 +491,8 @@ ActionExecutionValue executeAction(
boolean hasDiscoveredInputs)
throws ActionExecutionException, InterruptedException {
if (actionFileSystem != null) {
updateActionFileSystemContext(actionFileSystem, env, metadataHandler, expandedFilesets);
updateActionFileSystemContext(
action, actionFileSystem, env, metadataHandler, expandedFilesets);
}

ActionExecutionContext actionExecutionContext =
Expand Down Expand Up @@ -834,6 +837,7 @@ NestedSet<Artifact> discoverInputs(
threadStateReceiverFactory.apply(actionLookupData));
if (actionFileSystem != null) {
updateActionFileSystemContext(
action,
actionFileSystem,
env,
THROWING_METADATA_INJECTOR_FOR_ACTIONFS,
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -18,6 +18,7 @@
import com.google.common.collect.ImmutableList;
import com.google.common.collect.ImmutableMap;
import com.google.devtools.build.lib.actions.Action;
import com.google.devtools.build.lib.actions.ActionExecutionMetadata;
import com.google.devtools.build.lib.actions.ActionInputMap;
import com.google.devtools.build.lib.actions.Artifact;
import com.google.devtools.build.lib.actions.ArtifactPathResolver;
Expand Down Expand Up @@ -196,6 +197,7 @@ default FileSystem createActionFileSystem(
* @param filesets The Fileset symlinks known for this action.
*/
default void updateActionFileSystemContext(
ActionExecutionMetadata action,
FileSystem actionFileSystem,
Environment env,
MetadataInjector injector,
Expand Down
Loading

0 comments on commit 3b39ab8

Please sign in to comment.