Repository: pantsbuild/pants

Issue:
teach REAPI client to better respond to server backpressure and slowness
The remote cache client currently does not have any explicit rate limits or other way to respond to a slow server by backing off at a higher level than just retrying.

Several ideas to improve this situation (in no particular order):

* Add a configurable cap on total concurrent requests (possibly using existing tonic/tower filters)

* REAPI [specifies that clients should respect](https://github.com/bazelbuild/remote-apis/blob/3e385366f152e99adda5ab5e4857b1ab221ba2fe/build/bazel/remote/execution/v2/remote_execution.proto#L40-L43) a [`RetryInfo`](https://github.com/googleapis/googleapis/blob/f2c68a97c2f1344bc0809023c69c583a00a09b0b/google/rpc/error_details.proto#L27-L43) sent by REAPI server to set the minimum retry delay. This would give REAPI servers a way to force clients to back-off.

* Provide an explicit rate limit for remote store requests that can be set by the user. Consider whether rate limit should be number of requests, bandwidth cap, or both.

* Consider making the remote store rate limits dynamic based on observed server response latency and bandwidth.




Discussion:

Resolving patch:
diff --git a/src/python/pants/engine/internals/scheduler.py b/src/python/pants/engine/internals/scheduler.py
--- a/src/python/pants/engine/internals/scheduler.py
+++ b/src/python/pants/engine/internals/scheduler.py
@@ -175,14 +175,17 @@ def __init__(
             store_chunk_bytes=execution_options.remote_store_chunk_bytes,
             store_chunk_upload_timeout=execution_options.remote_store_chunk_upload_timeout_seconds,
             store_rpc_retries=execution_options.remote_store_rpc_retries,
+            store_rpc_concurrency=execution_options.remote_store_rpc_concurrency,
             cache_warnings_behavior=execution_options.remote_cache_warnings.value,
             cache_eager_fetch=execution_options.remote_cache_eager_fetch,
+            cache_rpc_concurrency=execution_options.remote_cache_rpc_concurrency,
             execution_extra_platform_properties=tuple(
                 tuple(pair.split("=", 1))
                 for pair in execution_options.remote_execution_extra_platform_properties
             ),
             execution_headers=tuple(execution_options.remote_execution_headers.items()),
             execution_overall_deadline_secs=execution_options.remote_execution_overall_deadline_secs,
+            execution_rpc_concurrency=execution_options.remote_execution_rpc_concurrency,
         )
         py_local_store_options = PyLocalStoreOptions(
             store_dir=local_store_options.store_dir,
diff --git a/src/python/pants/option/global_options.py b/src/python/pants/option/global_options.py
--- a/src/python/pants/option/global_options.py
+++ b/src/python/pants/option/global_options.py
@@ -149,6 +149,10 @@ class DynamicRemoteOptions:
     execution_address: str | None
     store_headers: dict[str, str]
     execution_headers: dict[str, str]
+    parallelism: int
+    store_rpc_concurrency: int
+    cache_rpc_concurrency: int
+    execution_rpc_concurrency: int
 
     @classmethod
     def disabled(cls) -> DynamicRemoteOptions:
@@ -161,6 +165,10 @@ def disabled(cls) -> DynamicRemoteOptions:
             execution_address=None,
             store_headers={},
             execution_headers={},
+            parallelism=DEFAULT_EXECUTION_OPTIONS.process_execution_remote_parallelism,
+            store_rpc_concurrency=DEFAULT_EXECUTION_OPTIONS.remote_store_rpc_concurrency,
+            cache_rpc_concurrency=DEFAULT_EXECUTION_OPTIONS.remote_cache_rpc_concurrency,
+            execution_rpc_concurrency=DEFAULT_EXECUTION_OPTIONS.remote_execution_rpc_concurrency,
         )
 
     @classmethod
@@ -180,6 +188,10 @@ def from_options(
         instance_name = cast("str | None", bootstrap_options.remote_instance_name)
         execution_headers = cast("dict[str, str]", bootstrap_options.remote_execution_headers)
         store_headers = cast("dict[str, str]", bootstrap_options.remote_store_headers)
+        parallelism = cast(int, bootstrap_options.process_execution_remote_parallelism)
+        store_rpc_concurrency = cast(int, bootstrap_options.remote_store_rpc_concurrency)
+        cache_rpc_concurrency = cast(int, bootstrap_options.remote_cache_rpc_concurrency)
+        execution_rpc_concurrency = cast(int, bootstrap_options.remote_execution_rpc_concurrency)
 
         if bootstrap_options.remote_oauth_bearer_token_path:
             oauth_token = (
@@ -289,6 +301,10 @@ def from_options(
             execution_address=execution_address,
             store_headers=store_headers,
             execution_headers=execution_headers,
+            parallelism=parallelism,
+            store_rpc_concurrency=store_rpc_concurrency,
+            cache_rpc_concurrency=cache_rpc_concurrency,
+            execution_rpc_concurrency=execution_rpc_concurrency,
         )
         return opts, auth_plugin_result
 
@@ -319,14 +335,17 @@ class ExecutionOptions:
     remote_store_chunk_bytes: Any
     remote_store_chunk_upload_timeout_seconds: int
     remote_store_rpc_retries: int
+    remote_store_rpc_concurrency: int
 
     remote_cache_eager_fetch: bool
     remote_cache_warnings: RemoteCacheWarningsBehavior
+    remote_cache_rpc_concurrency: int
 
     remote_execution_address: str | None
     remote_execution_extra_platform_properties: List[str]
     remote_execution_headers: Dict[str, str]
     remote_execution_overall_deadline_secs: int
+    remote_execution_rpc_concurrency: int
 
     @classmethod
     def from_options(
@@ -345,7 +364,7 @@ def from_options(
             # Process execution setup.
             process_execution_local_cache=bootstrap_options.process_execution_local_cache,
             process_execution_local_parallelism=bootstrap_options.process_execution_local_parallelism,
-            process_execution_remote_parallelism=bootstrap_options.process_execution_remote_parallelism,
+            process_execution_remote_parallelism=dynamic_remote_options.parallelism,
             process_execution_local_cleanup=bootstrap_options.process_execution_local_cleanup,
             process_execution_cache_namespace=bootstrap_options.process_execution_cache_namespace,
             # Remote store setup.
@@ -354,14 +373,17 @@ def from_options(
             remote_store_chunk_bytes=bootstrap_options.remote_store_chunk_bytes,
             remote_store_chunk_upload_timeout_seconds=bootstrap_options.remote_store_chunk_upload_timeout_seconds,
             remote_store_rpc_retries=bootstrap_options.remote_store_rpc_retries,
+            remote_store_rpc_concurrency=dynamic_remote_options.store_rpc_concurrency,
             # Remote cache setup.
             remote_cache_eager_fetch=bootstrap_options.remote_cache_eager_fetch,
             remote_cache_warnings=bootstrap_options.remote_cache_warnings,
+            remote_cache_rpc_concurrency=dynamic_remote_options.cache_rpc_concurrency,
             # Remote execution setup.
             remote_execution_address=dynamic_remote_options.execution_address,
             remote_execution_extra_platform_properties=bootstrap_options.remote_execution_extra_platform_properties,
             remote_execution_headers=dynamic_remote_options.execution_headers,
             remote_execution_overall_deadline_secs=bootstrap_options.remote_execution_overall_deadline_secs,
+            remote_execution_rpc_concurrency=dynamic_remote_options.execution_rpc_concurrency,
         )
 
 
@@ -428,9 +450,11 @@ def from_options(cls, options: OptionValueContainer) -> LocalStoreOptions:
     remote_store_chunk_bytes=1024 * 1024,
     remote_store_chunk_upload_timeout_seconds=60,
     remote_store_rpc_retries=2,
+    remote_store_rpc_concurrency=128,
     # Remote cache setup.
     remote_cache_eager_fetch=True,
     remote_cache_warnings=RemoteCacheWarningsBehavior.first_only,
+    remote_cache_rpc_concurrency=128,
     # Remote execution setup.
     remote_execution_address=None,
     remote_execution_extra_platform_properties=[],
@@ -438,6 +462,7 @@ def from_options(cls, options: OptionValueContainer) -> LocalStoreOptions:
         "user-agent": f"pants/{VERSION}",
     },
     remote_execution_overall_deadline_secs=60 * 60,  # one hour
+    remote_execution_rpc_concurrency=128,
 )
 
 DEFAULT_LOCAL_STORE_OPTIONS = LocalStoreOptions()
@@ -1159,6 +1184,13 @@ def register_bootstrap_options(cls, register):
             default=DEFAULT_EXECUTION_OPTIONS.remote_store_rpc_retries,
             help="Number of times to retry any RPC to the remote store before giving up.",
         )
+        register(
+            "--remote-store-rpc-concurrency",
+            type=int,
+            advanced=True,
+            default=DEFAULT_EXECUTION_OPTIONS.remote_store_rpc_concurrency,
+            help="The number of concurrent requests allowed to the remote store service.",
+        )
 
         register(
             "--remote-cache-warnings",
@@ -1182,6 +1214,13 @@ def register_bootstrap_options(cls, register):
                 "encountered by reducing the surface area of when remote caching is used."
             ),
         )
+        register(
+            "--remote-cache-rpc-concurrency",
+            type=int,
+            advanced=True,
+            default=DEFAULT_EXECUTION_OPTIONS.remote_cache_rpc_concurrency,
+            help="The number of concurrent requests allowed to the remote cache service.",
+        )
 
         register(
             "--remote-execution-address",
@@ -1221,6 +1260,13 @@ def register_bootstrap_options(cls, register):
             advanced=True,
             help="Overall timeout in seconds for each remote execution request from time of submission",
         )
+        register(
+            "--remote-execution-rpc-concurrency",
+            type=int,
+            advanced=True,
+            default=DEFAULT_EXECUTION_OPTIONS.remote_execution_rpc_concurrency,
+            help="The number of concurrent requests allowed to the remote execution service.",
+        )
 
         register(
             "--watch-filesystem",


Test patch:
diff --git a/src/rust/engine/fs/store/src/remote_tests.rs b/src/rust/engine/fs/store/src/remote_tests.rs
--- a/src/rust/engine/fs/store/src/remote_tests.rs
+++ b/src/rust/engine/fs/store/src/remote_tests.rs
@@ -157,6 +157,7 @@ async fn write_file_multiple_chunks() {
     10 * 1024,
     Duration::from_secs(5),
     1,
+    256,
   )
   .unwrap();
 
@@ -226,6 +227,7 @@ async fn write_connection_error() {
     10 * 1024 * 1024,
     Duration::from_secs(1),
     1,
+    256,
   )
   .unwrap();
   let error = store
@@ -301,6 +303,7 @@ fn new_byte_store(cas: &StubCAS) -> ByteStore {
     10 * MEGABYTES,
     Duration::from_secs(1),
     1,
+    256,
   )
   .unwrap()
 }
diff --git a/src/rust/engine/fs/store/src/tests.rs b/src/rust/engine/fs/store/src/tests.rs
--- a/src/rust/engine/fs/store/src/tests.rs
+++ b/src/rust/engine/fs/store/src/tests.rs
@@ -103,6 +103,7 @@ fn new_store<P: AsRef<Path>>(dir: P, cas_address: &str) -> Store {
       10 * MEGABYTES,
       Duration::from_secs(1),
       1,
+      256,
     )
     .unwrap()
 }
@@ -843,6 +844,7 @@ async fn instance_name_upload() {
       10 * MEGABYTES,
       Duration::from_secs(1),
       1,
+      256,
     )
     .unwrap();
 
@@ -870,6 +872,7 @@ async fn instance_name_download() {
       10 * MEGABYTES,
       Duration::from_secs(1),
       1,
+      256,
     )
     .unwrap();
 
@@ -919,6 +922,7 @@ async fn auth_upload() {
       10 * MEGABYTES,
       Duration::from_secs(1),
       1,
+      256,
     )
     .unwrap();
 
@@ -948,6 +952,7 @@ async fn auth_download() {
       10 * MEGABYTES,
       Duration::from_secs(1),
       1,
+      256,
     )
     .unwrap();
 
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
@@ -94,6 +94,7 @@ impl StoreSetup {
         10 * 1024 * 1024,
         Duration::from_secs(1),
         1,
+        256,
       )
       .unwrap();
     StoreSetup {
@@ -140,6 +141,7 @@ fn create_cached_runner(
       true,
       RemoteCacheWarningsBehavior::FirstOnly,
       eager_fetch,
+      256,
     )
     .expect("caching command runner"),
   );
@@ -587,6 +589,7 @@ async fn make_action_result_basic() {
     true,
     RemoteCacheWarningsBehavior::FirstOnly,
     false,
+    256,
   )
   .expect("caching command runner");
 
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
@@ -31,6 +31,9 @@ use tonic::{Code, Status};
 
 const OVERALL_DEADLINE_SECS: Duration = Duration::from_secs(10 * 60);
 const RETRY_INTERVAL: Duration = Duration::from_micros(0);
+const STORE_CONCURRENCY_LIMIT: usize = 256;
+const EXEC_CONCURRENCY_LIMIT: usize = 256;
+const CACHE_CONCURRENCY_LIMIT: usize = 256;
 
 #[derive(Debug, PartialEq)]
 pub(crate) struct RemoteTestResult {
@@ -866,6 +869,7 @@ async fn sends_headers() {
       10 * 1024 * 1024,
       Duration::from_secs(1),
       1,
+      STORE_CONCURRENCY_LIMIT,
     )
     .unwrap();
 
@@ -882,6 +886,8 @@ async fn sends_headers() {
     Platform::Linux,
     OVERALL_DEADLINE_SECS,
     RETRY_INTERVAL,
+    EXEC_CONCURRENCY_LIMIT,
+    CACHE_CONCURRENCY_LIMIT,
   )
   .unwrap();
   let context = Context {
@@ -1061,6 +1067,7 @@ async fn ensure_inline_stdio_is_stored() {
       10 * 1024 * 1024,
       Duration::from_secs(1),
       1,
+      STORE_CONCURRENCY_LIMIT,
     )
     .unwrap();
 
@@ -1074,6 +1081,8 @@ async fn ensure_inline_stdio_is_stored() {
     Platform::Linux,
     OVERALL_DEADLINE_SECS,
     RETRY_INTERVAL,
+    EXEC_CONCURRENCY_LIMIT,
+    CACHE_CONCURRENCY_LIMIT,
   )
   .unwrap();
 
@@ -1438,6 +1447,7 @@ async fn execute_missing_file_uploads_if_known() {
       10 * 1024 * 1024,
       Duration::from_secs(1),
       1,
+      STORE_CONCURRENCY_LIMIT,
     )
     .unwrap();
   store
@@ -1458,6 +1468,8 @@ async fn execute_missing_file_uploads_if_known() {
     Platform::Linux,
     OVERALL_DEADLINE_SECS,
     RETRY_INTERVAL,
+    EXEC_CONCURRENCY_LIMIT,
+    CACHE_CONCURRENCY_LIMIT,
   )
   .unwrap();
 
@@ -1513,6 +1525,7 @@ async fn execute_missing_file_errors_if_unknown() {
       10 * 1024 * 1024,
       Duration::from_secs(1),
       1,
+      STORE_CONCURRENCY_LIMIT,
     )
     .unwrap();
 
@@ -1526,6 +1539,8 @@ async fn execute_missing_file_errors_if_unknown() {
     Platform::Linux,
     OVERALL_DEADLINE_SECS,
     RETRY_INTERVAL,
+    EXEC_CONCURRENCY_LIMIT,
+    CACHE_CONCURRENCY_LIMIT,
   )
   .unwrap();
 
@@ -2225,6 +2240,8 @@ fn create_command_runner(
     platform,
     OVERALL_DEADLINE_SECS,
     RETRY_INTERVAL,
+    EXEC_CONCURRENCY_LIMIT,
+    CACHE_CONCURRENCY_LIMIT,
   )
   .expect("Failed to make command runner");
   (command_runner, store)
@@ -2274,6 +2291,7 @@ pub(crate) fn make_store(
       10 * 1024 * 1024,
       Duration::from_secs(1),
       1,
+      STORE_CONCURRENCY_LIMIT,
     )
     .unwrap()
 }