Repository: pantsbuild/pants

Issue:
Introduce an optional speculation delay for the remote cache
When `nailgun` is combined with a remote cache, cancellation of `nailgun` processes due to cache hits results in `nailgun` servers getting killed (as expected: see #14937 for more on that topic) (it can also currently cause an error: #16895 ... but that should be fixed independently).

But the cost of killing a `nailgun` server (which makes _future_ processes run more slowly) is not currently accounted for by our default setting of speculating for all requests. And figuring out the exact likelihood that we will hit the cache is a challenging problem.

Speculation is frequently (see [hedged requests](https://blog.acolyer.org/2015/01/15/the-tail-at-scale/)) configured with a delay before speculating: in this case, we should probably add a remote cache speculation delay which would be conditionally applied to processes which will use nailgun.

----

The core of the implementation for this would be to:
* Adjust the [remote_cache::CommandRunner](https://github.com/pantsbuild/pants/blob/f6b2832539e26ddffb19a7d665b9cef697abd0da/src/rust/engine/src/context.rs#L321-L338) to inspect [whether a process will use nailgun](https://github.com/pantsbuild/pants/blob/f6b2832539e26ddffb19a7d665b9cef697abd0da/src/rust/engine/src/context.rs#L226-L230), and introduce a delay before running the local process.
    * This would be implemented by adjusting the [`select!` call in the `remote_cache::CommandRunner`](https://github.com/pantsbuild/pants/blob/f6b2832539e26ddffb19a7d665b9cef697abd0da/src/rust/engine/process_execution/src/remote_cache.rs#L293-L327) to be inside of a `loop` which started by delaying for the delay period, and then if the delay period had elapsed, _or_ the `cache_read_future` had returned a miss, to run the local process.

How the configurable delay would be provided has two potential implementations though:
* Expose a field on `Process` like `remote_cache_speculation_delay`, and then use it from [JvmProcess](https://github.com/pantsbuild/pants/blob/f6b2832539e26ddffb19a7d665b9cef697abd0da/src/python/pants/jvm/jdk_rules.py#L308-L365) for processes which set `use_nailgun=True`. Allow the delay to be configured by the JVM subsystem, specific to this implementation.
* Add a global option next to [`--process-execution-enable-nailgun`](https://www.pantsbuild.org/docs/reference-global#process_execution_local_enable_nailgun) ([consumed here](https://github.com/pantsbuild/pants/blob/f6b2832539e26ddffb19a7d665b9cef697abd0da/src/rust/engine/src/context.rs#L203-L233)) which would apply to all nailgunned processes, rather than being configured on a `Process`-by-`Process` basis.

Between the two, having the value be configured on a `Process`-by-`Process` basis would provide the most flexibility for `@rule` authors.


Discussion:
once the result from the cache server is back and it is negative, how to cancel the sleep and run the process immediately?
Also, with remote_cache_speculation_delay per process, the command runner shouldn't care about nailgun or not, right? it would just check if the remote_cache_speculation_delay is non zero?
> once the result from the cache server is back and it is negative, how to cancel the sleep and run the process immediately?

Mm, yea: good point. This can't be a wrapping `CommandRunner`: it would need to actually be built in to the `remote_cache::CommandRunner`: in particular, [in this `select!` ](https://github.com/pantsbuild/pants/blob/f6b2832539e26ddffb19a7d665b9cef697abd0da/src/rust/engine/process_execution/src/remote_cache.rs#L293-L327). I'll update the description.

> Also, with remote_cache_speculation_delay per process, the command runner shouldn't care about nailgun or not, right? it would just check if the remote_cache_speculation_delay is non zero?

Yea, that's right.

Resolving patch:
diff --git a/src/python/pants/engine/process.py b/src/python/pants/engine/process.py
--- a/src/python/pants/engine/process.py
+++ b/src/python/pants/engine/process.py
@@ -67,6 +67,7 @@ class Process:
     execution_slot_variable: str | None
     concurrency_available: int
     cache_scope: ProcessCacheScope
+    remote_cache_speculation_delay_millis: int
 
     def __init__(
         self,
@@ -88,6 +89,7 @@ def __init__(
         concurrency_available: int = 0,
         cache_scope: ProcessCacheScope = ProcessCacheScope.SUCCESSFUL,
         platform: Platform | None = None,
+        remote_cache_speculation_delay_millis: int = 0,
     ) -> None:
         """Request to run a subprocess, similar to subprocess.Popen.
 
@@ -134,6 +136,7 @@ def __init__(
         self.execution_slot_variable = execution_slot_variable
         self.concurrency_available = concurrency_available
         self.cache_scope = cache_scope
+        self.remote_cache_speculation_delay_millis = remote_cache_speculation_delay_millis
 
         if platform is not None:
             warn_or_error(
diff --git a/src/python/pants/jvm/jdk_rules.py b/src/python/pants/jvm/jdk_rules.py
--- a/src/python/pants/jvm/jdk_rules.py
+++ b/src/python/pants/jvm/jdk_rules.py
@@ -323,6 +323,7 @@ class JvmProcess:
     extra_env: FrozenDict[str, str]
     cache_scope: ProcessCacheScope | None
     use_nailgun: bool
+    remote_cache_speculation_delay: int | None
 
     def __init__(
         self,
@@ -341,6 +342,7 @@ def __init__(
         timeout_seconds: int | float | None = None,
         cache_scope: ProcessCacheScope | None = None,
         use_nailgun: bool = True,
+        remote_cache_speculation_delay: int | None = None,
     ):
         self.jdk = jdk
         self.argv = tuple(argv)
@@ -357,6 +359,7 @@ def __init__(
         self.extra_immutable_input_digests = FrozenDict(extra_immutable_input_digests or {})
         self.extra_env = FrozenDict(extra_env or {})
         self.use_nailgun = use_nailgun
+        self.remote_cache_speculation_delay = remote_cache_speculation_delay
 
         if not use_nailgun and extra_nailgun_keys:
             raise AssertionError(
@@ -370,7 +373,7 @@ def __init__(
 
 @rule
 async def jvm_process(
-    bash: BashBinary, request: JvmProcess, global_options: GlobalOptions
+    bash: BashBinary, request: JvmProcess, jvm: JvmSubsystem, global_options: GlobalOptions
 ) -> Process:
 
     jdk = request.jdk
@@ -414,6 +417,12 @@ def valid_jvm_opt(opt: str) -> str:
     if request.use_nailgun:
         use_nailgun = [*jdk.immutable_input_digests, *request.extra_nailgun_keys]
 
+    remote_cache_speculation_delay_millis = 0
+    if request.remote_cache_speculation_delay is not None:
+        remote_cache_speculation_delay_millis = request.remote_cache_speculation_delay
+    elif request.use_nailgun:
+        remote_cache_speculation_delay_millis = jvm.nailgun_remote_cache_speculation_delay
+
     return Process(
         [*jdk.args(bash, request.classpath_entries), *jvm_options, *request.argv],
         input_digest=request.input_digest,
@@ -427,6 +436,7 @@ def valid_jvm_opt(opt: str) -> str:
         append_only_caches=jdk.append_only_caches,
         output_files=request.output_files,
         cache_scope=request.cache_scope or ProcessCacheScope.SUCCESSFUL,
+        remote_cache_speculation_delay_millis=remote_cache_speculation_delay_millis,
     )
 
 
diff --git a/src/python/pants/jvm/subsystems.py b/src/python/pants/jvm/subsystems.py
--- a/src/python/pants/jvm/subsystems.py
+++ b/src/python/pants/jvm/subsystems.py
@@ -3,7 +3,7 @@
 
 from __future__ import annotations
 
-from pants.option.option_types import BoolOption, DictOption, StrListOption, StrOption
+from pants.option.option_types import BoolOption, DictOption, IntOption, StrListOption, StrOption
 from pants.option.subsystem import Subsystem
 from pants.util.strutil import softwrap
 
@@ -98,3 +98,22 @@ class JvmSubsystem(Subsystem):
         ),
         advanced=True,
     )
+    # See https://github.com/pantsbuild/pants/issues/14937 for discussion of one way to improve
+    # our behavior around cancellation with nailgun.
+    nailgun_remote_cache_speculation_delay = IntOption(
+        default=1000,
+        help=softwrap(
+            """
+            The time in milliseconds to delay speculation of nailgun processes while reading
+            from the remote cache.
+
+            When speculating, a remote cache hit will cancel the local copy of a process. But
+            because nailgun does not natively support cancellation, that requires killing a
+            nailgun server, which will mean that future processes take longer to warm up.
+
+            This setting allows for trading off waiting for potentially slow cache entries
+            against potentially having to warm up a new nailgun server.
+            """
+        ),
+        advanced=True,
+    )


Test patch:
diff --git a/src/rust/engine/process_execution/src/remote_cache_tests.rs b/src/rust/engine/process_execution/src/remote_cache_tests.rs
--- a/src/rust/engine/process_execution/src/remote_cache_tests.rs
+++ b/src/rust/engine/process_execution/src/remote_cache_tests.rs
@@ -317,6 +317,7 @@ async fn cache_read_speculation() {
   async fn run_process(
     local_delay_ms: u64,
     remote_delay_ms: u64,
+    remote_cache_speculation_delay_ms: u64,
     cache_hit: bool,
     workunit: &mut RunningWorkunit,
   ) -> (i32, usize) {
@@ -330,6 +331,9 @@ async fn cache_read_speculation() {
       create_cached_runner(local_runner, &store_setup, CacheContentBehavior::Defer);
 
     let (process, action_digest) = create_process(&store_setup).await;
+    let process = process.remote_cache_speculation_delay(std::time::Duration::from_millis(
+      remote_cache_speculation_delay_ms,
+    ));
     if cache_hit {
       store_setup
         .cas
@@ -348,17 +352,32 @@ async fn cache_read_speculation() {
   }
 
   // Case 1: remote is faster than local.
-  let (exit_code, local_call_count) = run_process(200, 0, true, &mut workunit).await;
+  let (exit_code, local_call_count) = run_process(200, 0, 0, true, &mut workunit).await;
   assert_eq!(exit_code, 0);
   assert_eq!(local_call_count, 0);
 
   // Case 2: local is faster than remote.
-  let (exit_code, local_call_count) = run_process(0, 200, true, &mut workunit).await;
+  let (exit_code, local_call_count) = run_process(0, 200, 0, true, &mut workunit).await;
   assert_eq!(exit_code, 1);
   assert_eq!(local_call_count, 1);
 
   // Case 3: the remote lookup wins, but there is no cache entry so we fallback to local execution.
-  let (exit_code, local_call_count) = run_process(200, 0, false, &mut workunit).await;
+  let (exit_code, local_call_count) = run_process(200, 0, 0, false, &mut workunit).await;
+  assert_eq!(exit_code, 1);
+  assert_eq!(local_call_count, 1);
+
+  // Case 4: remote is faster than speculation read delay.
+  let (exit_code, local_call_count) = run_process(0, 100, 200, true, &mut workunit).await;
+  assert_eq!(exit_code, 0);
+  assert_eq!(local_call_count, 0);
+
+  // Case 5: remote is faster than speculation read delay, but there is no cache entry so we fallback to local execution.
+  let (exit_code, local_call_count) = run_process(0, 100, 200, false, &mut workunit).await;
+  assert_eq!(exit_code, 1);
+  assert_eq!(local_call_count, 1);
+
+  // Case 6: local with speculation read delay is faster than remote.
+  let (exit_code, local_call_count) = run_process(0, 200, 100, true, &mut workunit).await;
   assert_eq!(exit_code, 1);
   assert_eq!(local_call_count, 1);
 }
diff --git a/src/rust/engine/process_execution/src/remote_tests.rs b/src/rust/engine/process_execution/src/remote_tests.rs
--- a/src/rust/engine/process_execution/src/remote_tests.rs
+++ b/src/rust/engine/process_execution/src/remote_tests.rs
@@ -90,6 +90,7 @@ async fn make_execute_request() {
     concurrency_available: 0,
     cache_scope: ProcessCacheScope::Always,
     execution_strategy: ProcessExecutionStrategy::RemoteExecution(vec![]),
+    remote_cache_speculation_delay: std::time::Duration::from_millis(0),
   };
 
   let want_command = remexec::Command {
@@ -171,6 +172,7 @@ async fn make_execute_request_with_instance_name() {
       "target_platform".to_owned(),
       "apple-2e".to_owned(),
     )]),
+    remote_cache_speculation_delay: std::time::Duration::from_millis(0),
   };
 
   let want_command = remexec::Command {
@@ -255,6 +257,7 @@ async fn make_execute_request_with_cache_key_gen_version() {
     concurrency_available: 0,
     cache_scope: ProcessCacheScope::Always,
     execution_strategy: ProcessExecutionStrategy::RemoteExecution(vec![]),
+    remote_cache_speculation_delay: std::time::Duration::from_millis(0),
   };
 
   let mut want_command = remexec::Command {
@@ -482,6 +485,7 @@ async fn make_execute_request_with_timeout() {
     concurrency_available: 0,
     cache_scope: ProcessCacheScope::Always,
     execution_strategy: ProcessExecutionStrategy::RemoteExecution(vec![]),
+    remote_cache_speculation_delay: std::time::Duration::from_millis(0),
   };
 
   let want_command = remexec::Command {
@@ -588,6 +592,7 @@ async fn make_execute_request_using_immutable_inputs() {
     concurrency_available: 0,
     cache_scope: ProcessCacheScope::Always,
     execution_strategy: ProcessExecutionStrategy::RemoteExecution(vec![]),
+    remote_cache_speculation_delay: std::time::Duration::from_millis(0),
   };
 
   let want_command = remexec::Command {