chore: import upstream snapshot with attribution

mlflow/openai/api_request_parallel_processor.py

@@ -0,0 +1,131 @@
+# Based ons: https://github.com/openai/openai-cookbook/blob/6df6ceff470eeba26a56de131254e775292eac22/examples/api_request_parallel_processor.py
+# Several changes were made to make it work with MLflow.
+
+"""
+API REQUEST PARALLEL PROCESSOR
+
+Using the OpenAI API to process lots of text quickly takes some care.
+If you trickle in a million API requests one by one, they'll take days to complete.
+If you flood a million API requests in parallel, they'll exceed the rate limits and fail with
+errors. To maximize throughput, parallel requests need to be throttled to stay under rate limits.
+
+This script parallelizes requests to the OpenAI API
+
+Features:
+- Makes requests concurrently, to maximize throughput
+- Retries failed requests up to {max_attempts} times, to avoid missing data
+- Logs errors, to diagnose problems with requests
+"""
+
+from __future__ import annotations
+
+import logging
+import threading
+from concurrent.futures import FIRST_EXCEPTION, ThreadPoolExecutor, wait
+from dataclasses import dataclass
+from typing import Any, Callable
+
+import mlflow
+
+_logger = logging.getLogger(__name__)
+
+
+@dataclass
+class StatusTracker:
+    """Stores metadata about the script's progress. Only one instance is created."""
+
+    num_tasks_started: int = 0
+    num_tasks_in_progress: int = 0  # script ends when this reaches 0
+    num_tasks_succeeded: int = 0
+    num_tasks_failed: int = 0
+    num_rate_limit_errors: int = 0
+    lock: threading.Lock = threading.Lock()
+    error = None
+
+    def start_task(self):
+        with self.lock:
+            self.num_tasks_started += 1
+            self.num_tasks_in_progress += 1
+
+    def complete_task(self, *, success: bool):
+        with self.lock:
+            self.num_tasks_in_progress -= 1
+            if success:
+                self.num_tasks_succeeded += 1
+            else:
+                self.num_tasks_failed += 1
+
+    def increment_num_rate_limit_errors(self):
+        with self.lock:
+            self.num_rate_limit_errors += 1
+
+
+def call_api(
+    index: int,
+    results: list[tuple[int, Any]],
+    task: Callable[[], Any],
+    status_tracker: StatusTracker,
+):
+    import openai
+
+    status_tracker.start_task()
+    try:
+        result = task()
+        _logger.debug(f"Request #{index} succeeded")
+        status_tracker.complete_task(success=True)
+        results.append((index, result))
+    except openai.RateLimitError as e:
+        status_tracker.complete_task(success=False)
+        _logger.debug(f"Request #{index} failed with: {e}")
+        status_tracker.increment_num_rate_limit_errors()
+        status_tracker.error = mlflow.MlflowException(
+            f"Request #{index} failed with rate limit: {e}."
+        )
+    except Exception as e:
+        status_tracker.complete_task(success=False)
+        _logger.debug(f"Request #{index} failed with: {e}")
+        status_tracker.error = mlflow.MlflowException(
+            f"Request #{index} failed with: {e.__cause__}"
+        )
+
+
+def process_api_requests(
+    request_tasks: list[Callable[[], Any]],
+    max_workers: int = 10,
+):
+    """Processes API requests in parallel"""
+    # initialize trackers
+    status_tracker = StatusTracker()  # single instance to track a collection of variables
+
+    results: list[tuple[int, Any]] = []
+    request_tasks_iter = enumerate(request_tasks)
+    _logger.debug(f"Request pool executor will run {len(request_tasks)} requests")
+    with ThreadPoolExecutor(
+        max_workers=max_workers, thread_name_prefix="MlflowOpenAiApi"
+    ) as executor:
+        futures = [
+            executor.submit(
+                call_api,
+                index=index,
+                task=task,
+                results=results,
+                status_tracker=status_tracker,
+            )
+            for index, task in request_tasks_iter
+        ]
+        wait(futures, return_when=FIRST_EXCEPTION)
+
+    # after finishing, log final status
+    if status_tracker.num_tasks_failed > 0:
+        if status_tracker.num_tasks_failed == 1:
+            raise status_tracker.error
+        raise mlflow.MlflowException(
+            f"{status_tracker.num_tasks_failed} tasks failed. See logs for details."
+        )
+    if status_tracker.num_rate_limit_errors > 0:
+        _logger.debug(
+            f"{status_tracker.num_rate_limit_errors} rate limit errors received. "
+            "Consider running at a lower rate."
+        )
+
+    return [res for _, res in sorted(results)]