Skip to content

Commit 9610ae8

Browse files
alex-torokcopybara-github
authored andcommitted
Update PythonZipper action to use CommandLineItem.CapturingMapFn
### Summary This PR attempts to address bazelbuild#14890. This updates the PythonZipper action to use a `CommandLineItem.CapturingMapFn` to defer expanding arguments until after analysis. I used `CapturingMapFn`, as it seemed like the only way to go about implementing the arg filtering that is done by the current code to exclude the `executable` and `zipFile` artifacts. ### Test Plan Use the example from bazelbuild#14890 of looking at memory usage in tensorflow before/after this patch is applied on top of version `5.0.0` Initial Setup ``` git clone [email protected]:tensorflow/tensorflow.git cd tensorflow python3 -m venv venv source venv/bin/activate pip install numpy ``` View memory usage at 5.0.0: ```bash # STARTUP_FLAGS from https://bazel.build/rules/performance#enabling-memory-tracking $ bazel ${STARTUP_FLAGS[@]} build --nobuild //tensorflow/python/... $ bazel ${STARTUP_FLAGS[@]} dump --rules | egrep '(RULE|py_)' RULE COUNT ACTIONS BYTES EACH py_test 1,714 15,390 413,338,632 241,154 py_library 1,102 0 2,097,152 1,903 py_binary 33 198 8,914,840 270,146 py_runtime 6 0 0 0 py_runtime_pair 3 0 0 0 _concat_flatbuffer_py_srcs 2 2 0 0 $ bazel ${STARTUP_FLAGS[@]} info used-heap-size-after-gc 635MB ``` View memory usage at 5.0.0 with this patch applied: ```bash $ ~/code/bazel/bazel-bin/src/bazel ${STARTUP_FLAGS[@]} dump --rules | egrep '(RULE|py_)' RULE COUNT ACTIONS BYTES EACH py_test 1,714 15,390 65,323,312 38,111 py_library 1,102 0 2,359,576 2,141 py_binary 33 198 524,400 15,890 py_runtime 6 0 0 0 py_runtime_pair 3 0 0 0 _concat_flatbuffer_py_srcs 2 2 0 0 $ ~/code/bazel/bazel-bin/src/bazel ${STARTUP_FLAGS[@]} info used-heap-size-after-gc 283MB ``` Ensure that the generated actions have not changed: ```bash $ bazel aquery --build_python_zip 'mnemonic("PythonZipper", //tensorflow/python/...)' > /tmp/bazel_5_out $ ~/code/bazel/bazel-bin/src/bazel aquery --build_python_zip 'mnemonic("PythonZipper", //tensorflow/python/...)' > /tmp/bazel_5_patched_out $ diff /tmp/bazel_5_out /tmp/bazel_5_patched_out # note: no diff output ``` Closes bazelbuild#15037. PiperOrigin-RevId: 441257695
1 parent 6464f1c commit 9610ae8

File tree

2 files changed

+18
-4
lines changed

2 files changed

+18
-4
lines changed

src/main/java/com/google/devtools/build/lib/bazel/rules/python/BUILD

+1
Original file line numberDiff line numberDiff line change
@@ -22,6 +22,7 @@ java_library(
2222
deps = [
2323
"//src/main/java/com/google/devtools/build/lib/actions",
2424
"//src/main/java/com/google/devtools/build/lib/actions:artifacts",
25+
"//src/main/java/com/google/devtools/build/lib/actions:commandline_item",
2526
"//src/main/java/com/google/devtools/build/lib/analysis:actions/custom_command_line",
2627
"//src/main/java/com/google/devtools/build/lib/analysis:actions/launcher_file_write_action",
2728
"//src/main/java/com/google/devtools/build/lib/analysis:actions/substitution",

src/main/java/com/google/devtools/build/lib/bazel/rules/python/BazelPythonSemantics.java

+17-4
Original file line numberDiff line numberDiff line change
@@ -20,6 +20,7 @@
2020
import com.google.common.base.Preconditions;
2121
import com.google.common.collect.ImmutableList;
2222
import com.google.devtools.build.lib.actions.Artifact;
23+
import com.google.devtools.build.lib.actions.CommandLineItem;
2324
import com.google.devtools.build.lib.actions.ParamFileInfo;
2425
import com.google.devtools.build.lib.actions.ParameterFile;
2526
import com.google.devtools.build.lib.analysis.AnalysisUtils;
@@ -367,12 +368,24 @@ private static void createPythonZipAction(
367368

368369
// Read each runfile from execute path, add them into zip file at the right runfiles path.
369370
// Filter the executable file, cause we are building it.
371+
argv.addAll(
372+
CustomCommandLine.VectorArg.of(runfilesSupport.getRunfilesArtifacts())
373+
.mapped(
374+
(CommandLineItem.CapturingMapFn<Artifact>)
375+
(artifact, args) -> {
376+
if (!artifact.equals(executable) && !artifact.equals(zipFile)) {
377+
args.accept(
378+
getZipRunfilesPath(
379+
artifact.getRunfilesPath(),
380+
workspaceName,
381+
legacyExternalRunfiles)
382+
+ "="
383+
+ artifact.getExecPathString());
384+
}
385+
}));
386+
370387
for (Artifact artifact : runfilesSupport.getRunfilesArtifacts().toList()) {
371388
if (!artifact.equals(executable) && !artifact.equals(zipFile)) {
372-
argv.addDynamicString(
373-
getZipRunfilesPath(artifact.getRunfilesPath(), workspaceName, legacyExternalRunfiles)
374-
+ "="
375-
+ artifact.getExecPathString());
376389
inputsBuilder.add(artifact);
377390
}
378391
}

0 commit comments

Comments
 (0)