feat(opt-dist): new flag --benchmark-cargo-config

weihanglo · weihanglo · commit 0a11dcfdf4bd · 2024-06-03T17:15:52.000-04:00
The flag propagates cargo configs to `rustc-perf --cargo-config`,
which is particularly useful when the environment is air-gapped,
and you want to use the default set of training crates vendored
in the rustc-src tarball.
diff --git a/src/tools/opt-dist/src/environment.rs b/src/tools/opt-dist/src/environment.rs
@@ -17,6 +17,9 @@ pub struct Environment {
     host_llvm_dir: Utf8PathBuf,
     /// List of test paths that should be skipped when testing the optimized artifacts.
     skipped_tests: Vec<String>,
+    /// Arguments passed to `rustc-perf --cargo-config <value>` when running benchmarks.
+    #[builder(default)]
+    benchmark_cargo_config: Vec<String>,
     /// Directory containing a pre-built rustc-perf checkout.
     #[builder(default)]
     prebuilt_rustc_perf: Option<Utf8PathBuf>,
@@ -94,6 +97,10 @@ impl Environment {
     pub fn skipped_tests(&self) -> &[String] {
         &self.skipped_tests
     }
+
+    pub fn benchmark_cargo_config(&self) -> &[String] {
+        &self.benchmark_cargo_config
+    }
 }
 
 /// What is the extension of binary executables on this platform?
diff --git a/src/tools/opt-dist/src/main.rs b/src/tools/opt-dist/src/main.rs
@@ -90,6 +90,10 @@ enum EnvironmentCmd {
 
         #[clap(flatten)]
         shared: SharedArgs,
+
+        /// Arguments passed to `rustc-perf --cargo-config <value>` when running benchmarks.
+        #[arg(long)]
+        benchmark_cargo_config: Vec<String>,
     },
     /// Perform an optimized build on Linux CI, from inside Docker.
     LinuxCi {
@@ -119,6 +123,7 @@ fn create_environment(args: Args) -> anyhow::Result<(Environment, Vec<String>)>
             llvm_shared,
             use_bolt,
             skipped_tests,
+            benchmark_cargo_config,
             shared,
         } => {
             let env = EnvironmentBuilder::default()
@@ -132,6 +137,7 @@ fn create_environment(args: Args) -> anyhow::Result<(Environment, Vec<String>)>
                 .shared_llvm(llvm_shared)
                 .use_bolt(use_bolt)
                 .skipped_tests(skipped_tests)
+                .benchmark_cargo_config(benchmark_cargo_config)
                 .build()?;
 
             (env, shared.build_args)
diff --git a/src/tools/opt-dist/src/training.rs b/src/tools/opt-dist/src/training.rs
@@ -36,7 +36,7 @@ fn init_compiler_benchmarks(
     // Run rustc-perf benchmarks
     // Benchmark using profile_local with eprintln, which essentially just means
     // don't actually benchmark -- just make sure we run rustc a bunch of times.
-    cmd(&[
+    let mut cmd = cmd(&[
         env.cargo_stage_0().as_str(),
         "run",
         "-p",
@@ -61,7 +61,17 @@ fn init_compiler_benchmarks(
     .env("RUST_LOG", "collector=debug")
     .env("RUSTC", env.rustc_stage_0().as_str())
     .env("RUSTC_BOOTSTRAP", "1")
-    .workdir(&env.rustc_perf_dir())
+    .workdir(&env.rustc_perf_dir());
+
+    // This propagates cargo configs to `rustc-perf --cargo-config`,
+    // which is particularly useful when the environment is air-gapped,
+    // and you want to use the default set of training crates vendored
+    // in the rustc-src tarball.
+    for config in env.benchmark_cargo_config() {
+        cmd = cmd.arg("--cargo-config").arg(config);
+    }
+
+    cmd
 }
 
 /// Describes which `llvm-profdata` binary should be used for merging PGO profiles.
