chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,333 @@
|
||||
# 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.
|
||||
# Currently, only chat completion is supported.
|
||||
|
||||
"""
|
||||
API REQUEST PARALLEL PROCESSOR
|
||||
|
||||
Using the LangChain 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.
|
||||
This script parallelizes requests using LangChain API.
|
||||
|
||||
Features:
|
||||
- Streams requests from file, to avoid running out of memory for giant jobs
|
||||
- Makes requests concurrently, to maximize throughput
|
||||
- Logs errors, to diagnose problems with requests
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import logging
|
||||
import queue
|
||||
import threading
|
||||
import time
|
||||
import traceback
|
||||
from concurrent.futures import ThreadPoolExecutor
|
||||
from dataclasses import dataclass
|
||||
from typing import Any
|
||||
|
||||
from mlflow.langchain._compat import import_base_callback_handler, try_import_chain
|
||||
|
||||
BaseCallbackHandler = import_base_callback_handler()
|
||||
Chain = try_import_chain()
|
||||
|
||||
import mlflow
|
||||
from mlflow.exceptions import MlflowException
|
||||
from mlflow.langchain.utils.chat import (
|
||||
transform_request_json_for_chat_if_necessary,
|
||||
try_transform_response_iter_to_chat_format,
|
||||
try_transform_response_to_chat_format,
|
||||
)
|
||||
from mlflow.langchain.utils.serialization import convert_to_serializable
|
||||
from mlflow.pyfunc.context import Context, get_prediction_context
|
||||
from mlflow.tracing.utils import maybe_set_prediction_context
|
||||
|
||||
_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_api_errors: int = 0 # excluding rate limit errors, counted above
|
||||
lock: threading.Lock = threading.Lock()
|
||||
|
||||
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_api_errors(self):
|
||||
with self.lock:
|
||||
self.num_api_errors += 1
|
||||
|
||||
|
||||
@dataclass
|
||||
class APIRequest:
|
||||
"""
|
||||
Stores an API request's inputs, outputs, and other metadata. Contains a method to make an API
|
||||
call.
|
||||
|
||||
Args:
|
||||
index: The request's index in the tasks list
|
||||
lc_model: The LangChain model to call
|
||||
request_json: The request's input data
|
||||
results: The list to append the request's output data to, it's a list of tuples
|
||||
(index, response)
|
||||
errors: A dictionary to store any errors that occur
|
||||
convert_chat_responses: Whether to convert the model's responses to chat format
|
||||
did_perform_chat_conversion: Whether the input data was converted to chat format
|
||||
based on the model's type and input data.
|
||||
stream: Whether the request is a stream request
|
||||
prediction_context: The prediction context to use for the request
|
||||
"""
|
||||
|
||||
index: int
|
||||
lc_model: Any
|
||||
request_json: dict[str, Any]
|
||||
results: list[tuple[int, str]]
|
||||
errors: dict[int, str]
|
||||
convert_chat_responses: bool
|
||||
did_perform_chat_conversion: bool
|
||||
stream: bool
|
||||
params: dict[str, Any]
|
||||
prediction_context: Context | None = None
|
||||
|
||||
def _predict_single_input(self, single_input, callback_handlers, **kwargs):
|
||||
config = kwargs.pop("config", {})
|
||||
config["callbacks"] = config.get("callbacks", []) + (callback_handlers or [])
|
||||
if self.stream:
|
||||
return self.lc_model.stream(single_input, config=config, **kwargs)
|
||||
if hasattr(self.lc_model, "invoke"):
|
||||
return self.lc_model.invoke(single_input, config=config, **kwargs)
|
||||
else:
|
||||
# for backwards compatibility, __call__ is deprecated and will be removed in 0.3.0
|
||||
# kwargs shouldn't have config field if invoking with __call__
|
||||
return self.lc_model(single_input, callbacks=callback_handlers, **kwargs)
|
||||
|
||||
def _try_convert_response(self, response):
|
||||
if self.stream:
|
||||
return try_transform_response_iter_to_chat_format(response)
|
||||
else:
|
||||
return try_transform_response_to_chat_format(response)
|
||||
|
||||
def single_call_api(self, callback_handlers: list[BaseCallbackHandler] | None):
|
||||
from mlflow.langchain._compat import import_base_retriever
|
||||
from mlflow.langchain.utils.logging import langgraph_types, lc_runnables_types
|
||||
|
||||
BaseRetriever = import_base_retriever()
|
||||
|
||||
if isinstance(self.lc_model, BaseRetriever):
|
||||
# Retrievers are invoked differently than Chains
|
||||
response = self.lc_model.get_relevant_documents(
|
||||
**self.request_json, callbacks=callback_handlers, **self.params
|
||||
)
|
||||
elif isinstance(self.lc_model, lc_runnables_types() + langgraph_types()):
|
||||
if isinstance(self.request_json, dict):
|
||||
# This is a temporary fix for the case when spark_udf converts
|
||||
# input into pandas dataframe with column name, while the model
|
||||
# does not accept dictionaries as input, it leads to errors like
|
||||
# Expected Scalar value for String field 'query_text'
|
||||
try:
|
||||
response = self._predict_single_input(
|
||||
self.request_json, callback_handlers, **self.params
|
||||
)
|
||||
except TypeError as e:
|
||||
_logger.debug(
|
||||
f"Failed to invoke {self.lc_model.__class__.__name__} "
|
||||
f"with {self.request_json}. Error: {e!r}. Trying to "
|
||||
"invoke with the first value of the dictionary."
|
||||
)
|
||||
self.request_json = next(iter(self.request_json.values()))
|
||||
(
|
||||
prepared_request_json,
|
||||
did_perform_chat_conversion,
|
||||
) = transform_request_json_for_chat_if_necessary(
|
||||
self.request_json, self.lc_model
|
||||
)
|
||||
self.did_perform_chat_conversion = did_perform_chat_conversion
|
||||
|
||||
response = self._predict_single_input(
|
||||
prepared_request_json, callback_handlers, **self.params
|
||||
)
|
||||
else:
|
||||
response = self._predict_single_input(
|
||||
self.request_json, callback_handlers, **self.params
|
||||
)
|
||||
|
||||
if self.did_perform_chat_conversion or self.convert_chat_responses:
|
||||
response = self._try_convert_response(response)
|
||||
else:
|
||||
# return_only_outputs is invalid for stream call
|
||||
if Chain and isinstance(self.lc_model, Chain) and not self.stream:
|
||||
kwargs = {"return_only_outputs": True}
|
||||
else:
|
||||
kwargs = {}
|
||||
kwargs.update(**self.params)
|
||||
response = self._predict_single_input(self.request_json, callback_handlers, **kwargs)
|
||||
|
||||
if self.did_perform_chat_conversion or self.convert_chat_responses:
|
||||
response = self._try_convert_response(response)
|
||||
elif isinstance(response, dict) and len(response) == 1:
|
||||
# to maintain existing code, single output chains will still return
|
||||
# only the result
|
||||
response = response.popitem()[1]
|
||||
|
||||
return convert_to_serializable(response)
|
||||
|
||||
def call_api(
|
||||
self, status_tracker: StatusTracker, callback_handlers: list[BaseCallbackHandler] | None
|
||||
):
|
||||
"""
|
||||
Calls the LangChain API and stores results.
|
||||
"""
|
||||
_logger.debug(f"Request #{self.index} started with payload: {self.request_json}")
|
||||
|
||||
try:
|
||||
with maybe_set_prediction_context(self.prediction_context):
|
||||
response = self.single_call_api(callback_handlers)
|
||||
_logger.debug(f"Request #{self.index} succeeded with response: {response}")
|
||||
self.results.append((self.index, response))
|
||||
status_tracker.complete_task(success=True)
|
||||
except Exception as e:
|
||||
self.errors[self.index] = (
|
||||
f"error: {e!r} {traceback.format_exc()}\n request payload: {self.request_json}"
|
||||
)
|
||||
status_tracker.increment_num_api_errors()
|
||||
status_tracker.complete_task(success=False)
|
||||
|
||||
|
||||
def process_api_requests(
|
||||
lc_model,
|
||||
requests: list[Any | dict[str, Any]] | None = None,
|
||||
max_workers: int = 10,
|
||||
callback_handlers: list[BaseCallbackHandler] | None = None,
|
||||
convert_chat_responses: bool = False,
|
||||
params: dict[str, Any] | None = None,
|
||||
context: Context | None = None,
|
||||
):
|
||||
"""
|
||||
Processes API requests in parallel.
|
||||
"""
|
||||
|
||||
# initialize trackers
|
||||
retry_queue = queue.Queue()
|
||||
status_tracker = StatusTracker() # single instance to track a collection of variables
|
||||
next_request = None # variable to hold the next request to call
|
||||
context = context or get_prediction_context()
|
||||
|
||||
results = []
|
||||
errors = {}
|
||||
|
||||
# Note: we should call `transform_request_json_for_chat_if_necessary`
|
||||
# for the whole batch data, because the conversion should obey the rule
|
||||
# that if any record in the batch can't be converted, then all the record
|
||||
# in this batch can't be converted.
|
||||
(
|
||||
converted_chat_requests,
|
||||
did_perform_chat_conversion,
|
||||
) = transform_request_json_for_chat_if_necessary(requests, lc_model)
|
||||
|
||||
requests_iter = enumerate(converted_chat_requests)
|
||||
with ThreadPoolExecutor(
|
||||
max_workers=max_workers, thread_name_prefix="MlflowLangChainApi"
|
||||
) as executor:
|
||||
while True:
|
||||
# get next request (if one is not already waiting for capacity)
|
||||
if not retry_queue.empty():
|
||||
next_request = retry_queue.get_nowait()
|
||||
_logger.warning(f"Retrying request {next_request.index}: {next_request}")
|
||||
elif req := next(requests_iter, None):
|
||||
# get new request
|
||||
index, converted_chat_request_json = req
|
||||
next_request = APIRequest(
|
||||
index=index,
|
||||
lc_model=lc_model,
|
||||
request_json=converted_chat_request_json,
|
||||
results=results,
|
||||
errors=errors,
|
||||
convert_chat_responses=convert_chat_responses,
|
||||
did_perform_chat_conversion=did_perform_chat_conversion,
|
||||
stream=False,
|
||||
prediction_context=context,
|
||||
params=params,
|
||||
)
|
||||
status_tracker.start_task()
|
||||
else:
|
||||
next_request = None
|
||||
|
||||
# if enough capacity available, call API
|
||||
if next_request:
|
||||
# call API
|
||||
executor.submit(
|
||||
next_request.call_api,
|
||||
status_tracker=status_tracker,
|
||||
callback_handlers=callback_handlers,
|
||||
)
|
||||
|
||||
# if all tasks are finished, break
|
||||
# check next_request to avoid terminating the process
|
||||
# before extra requests need to be processed
|
||||
if status_tracker.num_tasks_in_progress == 0 and next_request is None:
|
||||
break
|
||||
|
||||
time.sleep(0.001) # avoid busy waiting
|
||||
|
||||
# after finishing, log final status
|
||||
if status_tracker.num_tasks_failed > 0:
|
||||
raise mlflow.MlflowException(
|
||||
f"{status_tracker.num_tasks_failed} tasks failed. Errors: {errors}"
|
||||
)
|
||||
|
||||
return [res for _, res in sorted(results)]
|
||||
|
||||
|
||||
def process_stream_request(
|
||||
lc_model,
|
||||
request_json: Any | dict[str, Any],
|
||||
callback_handlers: list[BaseCallbackHandler] | None = None,
|
||||
convert_chat_responses: bool = False,
|
||||
params: dict[str, Any] | None = None,
|
||||
):
|
||||
"""
|
||||
Process single stream request.
|
||||
"""
|
||||
if not hasattr(lc_model, "stream"):
|
||||
raise MlflowException(
|
||||
f"Model {lc_model.__class__.__name__} does not support streaming prediction output. "
|
||||
"No `stream` method found."
|
||||
)
|
||||
|
||||
(
|
||||
converted_chat_requests,
|
||||
did_perform_chat_conversion,
|
||||
) = transform_request_json_for_chat_if_necessary(request_json, lc_model)
|
||||
|
||||
api_request = APIRequest(
|
||||
index=0,
|
||||
lc_model=lc_model,
|
||||
request_json=converted_chat_requests,
|
||||
results=None,
|
||||
errors=None,
|
||||
convert_chat_responses=convert_chat_responses,
|
||||
did_perform_chat_conversion=did_perform_chat_conversion,
|
||||
stream=True,
|
||||
prediction_context=get_prediction_context(),
|
||||
params=params,
|
||||
)
|
||||
with maybe_set_prediction_context(api_request.prediction_context):
|
||||
return api_request.single_call_api(callback_handlers)
|
||||
Reference in New Issue
Block a user