wehub-resource-sync
132 lines
4.3 KiB
Python
132 lines
4.3 KiB
Python
# 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)]
|