feat: make rate limit a cmdarg

.cloud-build/execute_changed_notebooks_cli.py

@@ -121,6 +121,13 @@ parser.add_argument(
     default=True,
     help="Should run notebooks in parallel.",
 )
+parser.add_argument(
+    "--rate_limit",
+    type=int,
+    help="Number of parallel notebook executions per minute",
+    default=10,
+    required=False,
+)
 parser.add_argument(
     "--dry_run",
     type=str2bool,
@@ -168,5 +175,6 @@ else:
         variable_region=args.variable_region,
         variable_service_account=args.variable_service_account,
         variable_vpc_network=args.variable_vpc_network,
-        private_pool_id=args.private_pool_id
+        private_pool_id=args.private_pool_id,
+        rate_limit_count=args.rate_limit
 )

.cloud-build/execute_changed_notebooks_helper.py

@@ -234,7 +234,6 @@ def _create_tag(filepath: str) -> str:
     return tag
 
 
-rate_limit = RateLimit(max_count=10, per=60, greedy=False)
 
 
 def process_and_execute_notebook(
@@ -247,6 +246,7 @@ def process_and_execute_notebook(
     variable_vpc_network: Optional[str],
     private_pool_id: Optional[str],
     deadline: datetime.datetime,
+    rate_limit: int,
     notebook: str,
     should_get_tail_logs: bool = False,
 ) -> NotebookExecutionResult:
@@ -461,6 +461,7 @@ def process_and_execute_notebooks(
     variable_service_account: str,
     variable_vpc_network: Optional[str] = None,
     private_pool_id: Optional[str] = None,
+    rate_limit_count: Optional[int] = 10,
 ):
     """
     Run the notebooks that exist under the folders defined in the test_paths_file.
@@ -491,6 +492,7 @@ def process_and_execute_notebooks(
             Required. Should run notebooks in parallel using a thread pool as opposed to in sequence.
         timeout (str):
             Required. Timeout string according to https://cloud.google.com/build/docs/build-config-file-schema#timeout.
+        rate_limit_count (int): Max number of notebooks per minute to run in parallel.
     """
 
     # Calculate deadline
@@ -503,10 +505,13 @@ def process_and_execute_notebooks(
 
         print(f"Found {len(notebooks)} modified notebooks: {notebooks}")
 
+        rate_limit = RateLimit(max_count=rate_limit_count, per=60, greedy=False)
         if should_parallelize and len(notebooks) > 1:
             print(
                 "Running notebooks in parallel, so no logs will be displayed. Please wait..."
             )
+
+
             with concurrent.futures.ThreadPoolExecutor(max_workers=100) as executor:
                 print(f"Max workers: {executor._max_workers}")
 
@@ -523,6 +528,7 @@ def process_and_execute_notebooks(
                             variable_vpc_network,
                             private_pool_id,
                             deadline,
+                            rate_limit,
                         ),
                         notebooks,
                     )
@@ -540,6 +546,7 @@ def process_and_execute_notebooks(
                     private_pool_id=private_pool_id,
                     deadline=deadline,
                     notebook=notebook,
+                    rate_limit=rate_limit,
                 )
                 for notebook in notebooks
             ]