tune: try thread safe rate limiter

2026-09-26 14:42:04 +00:00 · 2023-04-11 18:34:16 +00:00
2 changed files with 7 additions and 4 deletions
@@ -32,7 +32,8 @@ import execute_notebook_helper
 import execute_notebook_remote
 import nbformat
 from google.cloud.devtools.cloudbuild_v1.types import BuildOperationMetadata
-from ratemate import RateLimit
+from pyrate_limiter import (Duration, RequestRate,
+                            Limiter)
 from tabulate import tabulate
 from utils import NotebookProcessors, util

@@ -156,9 +157,11 @@ def _create_tag(filepath: str) -> str:
    return tag


-rate_limit = RateLimit(max_count=10, per=60, greedy=True)
+rate_limit = RequestRate(10, Duration.HOUR)
+limiter = Limiter(rate_limit)


+@limiter.ratelimit('testing', delay=True, max_delay=360)
 def process_and_execute_notebook(
    container_uri: str,
    staging_bucket: str,
@@ -172,7 +175,6 @@ def process_and_execute_notebook(
    notebook: str,
    should_get_tail_logs: bool = False,
 ) -> NotebookExecutionResult:
-    rate_limit.wait()  # wait before creating the task

    print(f"Running notebook: {notebook}")

@@ -11,4 +11,5 @@ google-cloud-storage
 google-cloud-build
 ratemate
 GitPython
-tqdm
+tqdm
+pyrate-limiter