Add remote persistent worker support

Differential Revision: D65214397

Author: KapJI

.github/actions/build_example_persistent_worker/action.yml

@@ -0,0 +1,16 @@
+name: build_example_persistent_worker
+inputs:
+  buildbuddyApiKey:
+    description: "The API key for BuildBuddy remote cache and execution."
+    required: true
+runs:
+  using: composite
+  steps:
+  - name: Build examples/persistent_worker directory
+    env:
+      BUILDBUDDY_API_KEY: ${{ inputs.buildbuddyApiKey }}
+    run: |-
+      cd examples/persistent_worker
+      export PATH="$RUNNER_TEMP/artifacts:$PATH"
+      ./test.sh
+    shell: bash

.github/workflows/build-and-test.yml

@@ -2,6 +2,7 @@ name: Build and test
 on:
   push:
   pull_request:
+  workflow_dispatch: # allows manual triggering
 jobs:
   linux-build-and-test:
     runs-on: 4-core-ubuntu
@@ -69,6 +70,9 @@ jobs:
         $RUNNER_TEMP/artifacts/buck2 test //... -v 2
     - uses: ./.github/actions/build_example_conan
     - uses: ./.github/actions/build_example_no_prelude
+    - uses: ./.github/actions/build_example_persistent_worker
+      with:
+        buildbuddyApiKey: ${{ secrets.BUILDBUDDY_API_KEY }}
     - uses: ./.github/actions/setup_reindeer
     - uses: ./.github/actions/build_bootstrap
   windows-build-examples:

app/buck2_action_impl/src/actions/impls/run.rs

@@ -255,6 +255,7 @@ struct UnpackedWorkerValues<'v> {
     exe: &'v dyn CommandLineArgLike,
     id: WorkerId,
     concurrency: Option<usize>,
+    remote: bool,
 }
 
 struct UnpackedRunActionValues<'v> {
@@ -310,6 +311,7 @@ impl RunAction {
             exe: worker.exe_command_line(),
             id: WorkerId(worker.id),
             concurrency: worker.concurrency(),
+            remote: worker.remote(),
         });
 
         Ok(UnpackedRunActionValues {
@@ -325,6 +327,7 @@ impl RunAction {
         &self,
         fs: &ExecutorFs,
         artifact_visitor: &mut impl CommandLineArtifactVisitor,
+        actx: &dyn ActionExecutionCtx,
     ) -> anyhow::Result<(ExpandedCommandLine, Option<WorkerSpec>)> {
         let mut ctx = DefaultCommandLineContext::new(fs);
         let values = Self::unpack(&self.starlark_values)?;
@@ -341,10 +344,27 @@ impl RunAction {
                 .exe
                 .add_to_command_line(&mut worker_rendered, &mut ctx)?;
             worker.exe.visit_artifacts(artifact_visitor)?;
+            let worker_key = if worker.remote {
+                let mut worker_visitor = SimpleCommandLineArtifactVisitor::new();
+                worker.exe.visit_artifacts(&mut worker_visitor)?;
+                if !worker_visitor.outputs.is_empty() {
+                    // TODO[AH] create appropriate error enum value.
+                    return Err(anyhow::anyhow!("remote persistent worker command should not produce an output"));
+                }
+                let worker_inputs: Vec<&ArtifactGroupValues> = worker_visitor
+                    .inputs()
+                    .map(|group| actx.artifact_values(group))
+                    .collect();
+                let (_, worker_digest) = metadata_content(fs.fs(), &worker_inputs, actx.digest_config())?;
+                Some(worker_digest)
+            } else {
+                None
+            };
             Some(WorkerSpec {
                 exe: worker_rendered,
                 id: worker.id,
                 concurrency: worker.concurrency,
+                remote_key: worker_key,
             })
         } else {
             None
@@ -416,7 +436,7 @@ impl RunAction {
         let fs = executor_fs.fs();
 
         let (expanded, worker) =
-            self.expand_command_line_and_worker(&ctx.executor_fs(), visitor)?;
+            self.expand_command_line_and_worker(&ctx.executor_fs(), visitor, ctx)?;
 
         // TODO (@torozco): At this point, might as well just receive the list already. Finding
         // those things in a HashMap is just not very useful.

app/buck2_build_api/src/actions/execute/action_executor.rs

@@ -766,6 +766,7 @@ mod tests {
                 CommandGenerationOptions {
                     path_separator: PathSeparatorKind::Unix,
                     output_paths_behavior: Default::default(),
+                    use_remote_persistent_workers: false,
                 },
                 Default::default(),
             ),

app/buck2_build_api/src/interpreter/rule_defs/command_executor_config.rs

@@ -82,6 +82,7 @@ pub fn register_command_executor_config(builder: &mut GlobalsBuilder) {
     /// * `allow_hybrid_fallbacks_on_failure`: Whether to allow fallbacks when the result is failure (i.e. the command failed on the primary, but the infra worked)
     /// * `use_windows_path_separators`: Whether to use Windows path separators in command line arguments
     /// * `use_persistent workers`: Whether to use persistent workers for local execution if they are available
+    /// * `use_remote_persistent_workers`: Whether to use persistent workers for remote execution if they are available
     /// * `allow_cache_uploads`: Whether to upload local actions to the RE cache
     /// * `max_cache_upload_mebibytes`: Maximum size to upload in cache uploads
     /// * `experimental_low_pass_filter`: Whether to use the experimental low pass filter
@@ -106,6 +107,7 @@ pub fn register_command_executor_config(builder: &mut GlobalsBuilder) {
         #[starlark(default = false, require = named)] allow_hybrid_fallbacks_on_failure: bool,
         #[starlark(default = false, require = named)] use_windows_path_separators: bool,
         #[starlark(default = false, require = named)] use_persistent_workers: bool,
+        #[starlark(default = false, require = named)] use_remote_persistent_workers: bool,
         #[starlark(default = false, require = named)] allow_cache_uploads: bool,
         #[starlark(default = NoneOr::None, require = named)] max_cache_upload_mebibytes: NoneOr<
             i32,
@@ -307,6 +309,7 @@ pub fn register_command_executor_config(builder: &mut GlobalsBuilder) {
                         PathSeparatorKind::Unix
                     },
                     output_paths_behavior,
+                    use_remote_persistent_workers,
                 },
             }
         };

app/buck2_build_api/src/interpreter/rule_defs/provider/builtin/worker_info.rs

@@ -45,6 +45,8 @@ pub struct WorkerInfoGen<V: ValueLifetimeless> {
     pub exe: ValueOfUncheckedGeneric<V, FrozenStarlarkCmdArgs>,
     // Maximum number of concurrent commands to execute on a worker instance without queuing
     pub concurrency: ValueOfUncheckedGeneric<V, NoneOr<usize>>,
+    // Remote execution capable worker
+    pub remote: ValueOfUncheckedGeneric<V, bool>,
 
     pub id: u64,
 }
@@ -62,6 +64,7 @@ fn worker_info_creator(globals: &mut GlobalsBuilder) {
         #[starlark(require = named, default = NoneOr::None)] concurrency: NoneOr<
             ValueOf<'v, usize>,
         >,
+        #[starlark(require = named, default = false)] remote: bool,
         eval: &mut Evaluator<'v, '_, '_>,
     ) -> anyhow::Result<WorkerInfo<'v>> {
         let heap = eval.heap();
@@ -72,6 +75,7 @@ fn worker_info_creator(globals: &mut GlobalsBuilder) {
             exe,
             id,
             concurrency: heap.alloc_typed_unchecked(concurrency).cast(),
+            remote: heap.alloc_typed_unchecked(remote).cast(),
         })
     }
 }
@@ -90,6 +94,13 @@ impl<'v, V: ValueLike<'v>> WorkerInfoGen<V> {
             .expect("validated at construction")
             .into_option()
     }
+
+    pub fn remote(&self) -> bool {
+        self.remote
+            .to_value()
+            .unpack()
+            .expect("validated at construction")
+    }
 }
 
 fn validate_worker_info<'v, V>(info: &WorkerInfoGen<V>) -> anyhow::Result<()>

app/buck2_build_api_tests/src/interpreter/rule_defs/provider/builtin/worker_info.rs

@@ -24,7 +24,7 @@ fn run_display() {
         .run_starlark_bzl_test(
             r#"
 def test():
-    assert_eq('WorkerInfo(exe=cmd_args("x"), concurrency=None)', str(WorkerInfo(exe="x")))
+    assert_eq('WorkerInfo(exe=cmd_args("x"), concurrency=None, remote=False)', str(WorkerInfo(exe="x")))
 "#,
         )
         .unwrap();

app/buck2_core/src/execution_types/executor_config.rs

@@ -263,6 +263,7 @@ impl Default for CacheUploadBehavior {
 pub struct CommandGenerationOptions {
     pub path_separator: PathSeparatorKind,
     pub output_paths_behavior: OutputPathsBehavior,
+    pub use_remote_persistent_workers: bool,
 }
 
 #[derive(Debug, Eq, PartialEq, Hash, Allocative, Clone)]
@@ -294,6 +295,7 @@ impl CommandExecutorConfig {
             options: CommandGenerationOptions {
                 path_separator: PathSeparatorKind::system_default(),
                 output_paths_behavior: Default::default(),
+                use_remote_persistent_workers: false,
             },
         })
     }

app/buck2_execute/src/execute/command_executor.rs

@@ -11,6 +11,7 @@ use std::ops::ControlFlow;
 use std::sync::Arc;
 use std::time::Duration;
 
+use anyhow::anyhow;
 use anyhow::Context;
 use buck2_common::file_ops::TrackedFileDigest;
 use buck2_core::execution_types::executor_config::CommandGenerationOptions;
@@ -184,15 +185,31 @@ impl CommandExecutor {
                 }
                 CommandExecutionInput::ScratchPath(_) => None,
             });
+            let mut platform = self.0.re_platform.clone();
+            let args = if self.0.options.use_remote_persistent_workers && let Some(worker) = request.worker() && let Some(key) = worker.remote_key.as_ref() {
+                platform.properties.push(RE::Property {
+                    name: "persistentWorkerKey".to_owned(),
+                    value: key.to_string(),
+                });
+                // TODO[AH] Ideally, Buck2 could generate an argfile on the fly.
+                for arg in request.args() {
+                    if !(arg.starts_with("@") || arg.starts_with("-flagfile") || arg.starts_with("--flagfile")) {
+                        return Err(anyhow!("Remote persistent worker arguments must be passed as `@argfile`, `-flagfile=argfile`, or `--flagfile=argfile`."));
+                    }
+                }
+                worker.exe.iter().chain(request.args().iter()).cloned().collect()
+            } else {
+                request.all_args_vec()
+            };
             let action = re_create_action(
-                request.all_args_vec(),
+                args,
                 request.paths().output_paths(),
                 request.working_directory(),
                 request.env(),
                 input_digest,
                 action_metadata_blobs,
                 request.timeout(),
-                self.0.re_platform.clone(),
+                platform,
                 false,
                 digest_config,
                 self.0.options.output_paths_behavior,

app/buck2_execute/src/execute/request.rs

@@ -276,11 +276,12 @@ impl CommandExecutionPaths {
 #[derive(Copy, Clone, Dupe, Debug, Display, Allocative, Hash, PartialEq, Eq)]
 pub struct WorkerId(pub u64);
 
-#[derive(Clone)]
+#[derive(Clone, Debug)]
 pub struct WorkerSpec {
     pub id: WorkerId,
     pub exe: Vec<String>,
     pub concurrency: Option<usize>,
+    pub remote_key: Option<TrackedFileDigest>,
 }
 
 /// The data contains the information about the command to be executed.

app/buck2_execute/src/lib.rs

@@ -14,6 +14,7 @@
 #![feature(try_trait_v2)]
 #![feature(used_with_arg)]
 #![feature(trait_upcasting)]
+#![feature(let_chains)]
 
 pub mod artifact;
 pub mod artifact_utils;

app/buck2_server/src/daemon/common.rs

@@ -486,6 +486,7 @@ pub fn get_default_executor_config(host_platform: HostPlatformOverride) -> Comma
         options: CommandGenerationOptions {
             path_separator: get_default_path_separator(host_platform),
             output_paths_behavior: Default::default(),
+            use_remote_persistent_workers: false,
         },
     }
 }

app/buck2_test/src/orchestrator.rs

@@ -1171,6 +1171,7 @@ impl<'b> BuckTestOrchestrator<'b> {
             options: CommandGenerationOptions {
                 path_separator: PathSeparatorKind::system_default(),
                 output_paths_behavior: Default::default(),
+                use_remote_persistent_workers: false,
             },
         };
         let CommandExecutorResponse {
@@ -1698,6 +1699,7 @@ impl<'a> Execute2RequestExpander<'a> {
                     exe: worker_rendered,
                     id: WorkerId(worker.id),
                     concurrency: worker.concurrency(),
+                    remote_key: None,
                 })
             }
             _ => None,

examples/persistent_worker/.buckconfig

@@ -0,0 +1,17 @@
+[cells]
+root = .
+prelude = prelude
+toolchains = toolchains
+none = none
+
+[cell_aliases]
+config = prelude
+fbcode = none
+fbsource = none
+buck = none
+
+[external_cells]
+  prelude = bundled
+
+[parser]
+target_platform_detector_spec = target:root//...->prelude//platforms:default

examples/persistent_worker/.buckconfig.buildbuddy

@@ -0,0 +1,13 @@
+[buck2]
+digest_algorithms = SHA256
+
+[buck2_re_client]
+engine_address       = grpc://remote.buildbuddy.io
+action_cache_address = grpc://remote.buildbuddy.io
+cas_address          = grpc://remote.buildbuddy.io
+tls                  = true
+http_headers         = \
+  x-buildbuddy-api-key:$BUILDBUDDY_API_KEY
+
+[build]
+execution_platforms = root//platforms:buildbuddy

examples/persistent_worker/.buckconfig.buildbuddy-persistent-workers

@@ -0,0 +1,13 @@
+[buck2]
+digest_algorithms = SHA256
+
+[buck2_re_client]
+engine_address       = grpc://remote.buildbuddy.io
+action_cache_address = grpc://remote.buildbuddy.io
+cas_address          = grpc://remote.buildbuddy.io
+tls                  = true
+http_headers         = \
+  x-buildbuddy-api-key:$BUILDBUDDY_API_KEY
+
+[build]
+execution_platforms = root//platforms:buildbuddy-persistent-workers

examples/persistent_worker/.buckconfig.local-persistent-workers

@@ -0,0 +1,2 @@
+[build]
+execution_platforms = root//platforms:local-persistent-workers

examples/persistent_worker/.buckconfig.no-workers

@@ -0,0 +1,2 @@
+[build]
+execution_platforms = root//platforms:local

examples/persistent_worker/.buckroot

@@ -0,0 +1,3 @@
+# specify the following:
+# - BUILDBUDDY_API_KEY
+source_env_if_exists .envrc.private

examples/persistent_worker/.envrc

@@ -0,0 +1,4 @@
+.buckconfig.local
+.direnv
+.envrc.private
+prelude

examples/persistent_worker/.gitignore

@@ -0,0 +1,26 @@
+load("defs.bzl", "demo", "worker")
+
+python_binary(
+    name = "one_shot",
+    main = "one_shot.py",
+)
+
+python_binary(
+    name = "worker_py",
+    main = "persistent_worker.py",
+    deps = [
+        "//proto/bazel:worker_protocol_pb2",
+        "//proto/buck2:worker_pb2",
+    ],
+)
+
+worker(
+    name = "worker",
+    worker = ":worker_py",
+    visibility = ["PUBLIC"],
+)
+
+[
+    demo(name = "demo-" + str(i))
+    for i in range(4)
+]