wehub-resource-sync
367 lines
14 KiB
Python
367 lines
14 KiB
Python
"""
|
|
Defines an AsyncLoggingQueue that provides async fashion logging of metrics/tags/params using
|
|
queue based approach.
|
|
"""
|
|
|
|
import atexit
|
|
import enum
|
|
import logging
|
|
import threading
|
|
from concurrent.futures import ThreadPoolExecutor
|
|
from queue import Empty, Queue
|
|
from typing import Callable
|
|
|
|
from mlflow.entities.metric import Metric
|
|
from mlflow.entities.param import Param
|
|
from mlflow.entities.run_tag import RunTag
|
|
from mlflow.environment_variables import (
|
|
MLFLOW_ASYNC_LOGGING_BUFFERING_SECONDS,
|
|
MLFLOW_ASYNC_LOGGING_THREADPOOL_SIZE,
|
|
)
|
|
from mlflow.utils.async_logging.run_batch import RunBatch
|
|
from mlflow.utils.async_logging.run_operations import RunOperations
|
|
|
|
_logger = logging.getLogger(__name__)
|
|
|
|
|
|
ASYNC_LOGGING_WORKER_THREAD_PREFIX = "MLflowBatchLoggingWorkerPool"
|
|
ASYNC_LOGGING_STATUS_CHECK_THREAD_PREFIX = "MLflowAsyncLoggingStatusCheck"
|
|
|
|
|
|
class QueueStatus(enum.Enum):
|
|
"""Status of the async queue"""
|
|
|
|
# The queue is listening to new data and logging enqueued data to MLflow.
|
|
ACTIVE = 1
|
|
# The queue is not listening to new data, but still logging enqueued data to MLflow.
|
|
TEAR_DOWN = 2
|
|
# The queue is neither listening to new data or logging enqueued data to MLflow.
|
|
IDLE = 3
|
|
|
|
|
|
_MAX_ITEMS_PER_BATCH = 1000
|
|
_MAX_PARAMS_PER_BATCH = 100
|
|
_MAX_TAGS_PER_BATCH = 100
|
|
|
|
|
|
class AsyncLoggingQueue:
|
|
"""
|
|
This is a queue based run data processor that queues incoming batches and processes them using
|
|
single worker thread.
|
|
"""
|
|
|
|
def __init__(
|
|
self, logging_func: Callable[[str, list[Metric], list[Param], list[RunTag]], None]
|
|
) -> None:
|
|
"""Initializes an AsyncLoggingQueue object.
|
|
|
|
Args:
|
|
logging_func: A callable function that takes in four arguments: a string
|
|
representing the run_id, a list of Metric objects,
|
|
a list of Param objects, and a list of RunTag objects.
|
|
"""
|
|
self._queue = Queue()
|
|
self._lock = threading.RLock()
|
|
self._logging_func = logging_func
|
|
|
|
self._stop_data_logging_thread_event = threading.Event()
|
|
self._status = QueueStatus.IDLE
|
|
|
|
def _at_exit_callback(self) -> None:
|
|
"""Callback function to be executed when the program is exiting.
|
|
|
|
Stops the data processing thread and waits for the queue to be drained. Finally, shuts down
|
|
the thread pools used for data logging and batch processing status check.
|
|
"""
|
|
try:
|
|
# Stop the data processing thread
|
|
self._stop_data_logging_thread_event.set()
|
|
# Waits till logging queue is drained.
|
|
self._batch_logging_thread.join()
|
|
self._batch_logging_worker_threadpool.shutdown(wait=True)
|
|
self._batch_status_check_threadpool.shutdown(wait=True)
|
|
except Exception as e:
|
|
_logger.error(f"Encountered error while trying to finish logging: {e}")
|
|
|
|
def end_async_logging(self) -> None:
|
|
with self._lock:
|
|
# Stop the data processing thread.
|
|
self._stop_data_logging_thread_event.set()
|
|
# Waits till logging queue is drained.
|
|
self._batch_logging_thread.join()
|
|
# Set the status to tear down. The worker threads will still process
|
|
# the remaining data.
|
|
self._status = QueueStatus.TEAR_DOWN
|
|
# Clear the status to avoid blocking next logging.
|
|
self._stop_data_logging_thread_event.clear()
|
|
|
|
def shut_down_async_logging(self) -> None:
|
|
"""
|
|
Shut down the async logging queue and wait for the queue to be drained.
|
|
Use this method if the async logging should be terminated.
|
|
"""
|
|
self.end_async_logging()
|
|
self._batch_logging_worker_threadpool.shutdown(wait=True)
|
|
self._batch_status_check_threadpool.shutdown(wait=True)
|
|
self._status = QueueStatus.IDLE
|
|
|
|
def flush(self) -> None:
|
|
"""
|
|
Flush the async logging queue and restart thread to listen
|
|
to incoming data after flushing.
|
|
|
|
Calling this method will flush the queue to ensure all the data are logged.
|
|
"""
|
|
self.shut_down_async_logging()
|
|
# Reinitialize the logging thread and set the status to active.
|
|
self.activate()
|
|
|
|
def _logging_loop(self) -> None:
|
|
"""
|
|
Continuously logs run data until `self._continue_to_process_data` is set to False.
|
|
If an exception occurs during logging, a `MlflowException` is raised.
|
|
"""
|
|
try:
|
|
while not self._stop_data_logging_thread_event.is_set():
|
|
self._log_run_data()
|
|
# Drain the queue after the stop event is set.
|
|
while not self._queue.empty():
|
|
self._log_run_data()
|
|
except Exception as e:
|
|
from mlflow.exceptions import MlflowException
|
|
|
|
raise MlflowException(f"Exception inside the run data logging thread: {e}")
|
|
|
|
def _fetch_batch_from_queue(self) -> list[RunBatch]:
|
|
"""Fetches a batch of run data from the queue.
|
|
|
|
Returns:
|
|
RunBatch: A batch of run data.
|
|
"""
|
|
batches = []
|
|
if self._queue.empty():
|
|
return batches
|
|
queue_size = self._queue.qsize() # Estimate the queue's size.
|
|
merged_batch = self._queue.get()
|
|
for i in range(queue_size - 1):
|
|
if self._queue.empty():
|
|
# `queue_size` is an estimate, so we need to check if the queue is empty.
|
|
break
|
|
batch = self._queue.get()
|
|
|
|
if (
|
|
merged_batch.run_id != batch.run_id
|
|
or (
|
|
len(merged_batch.metrics + merged_batch.params + merged_batch.tags)
|
|
+ len(batch.metrics + batch.params + batch.tags)
|
|
)
|
|
>= _MAX_ITEMS_PER_BATCH
|
|
or len(merged_batch.params) + len(batch.params) >= _MAX_PARAMS_PER_BATCH
|
|
or len(merged_batch.tags) + len(batch.tags) >= _MAX_TAGS_PER_BATCH
|
|
):
|
|
# Make a new batch if the run_id is different or the batch is full.
|
|
batches.append(merged_batch)
|
|
merged_batch = batch
|
|
else:
|
|
merged_batch.add_child_batch(batch)
|
|
merged_batch.params.extend(batch.params)
|
|
merged_batch.tags.extend(batch.tags)
|
|
merged_batch.metrics.extend(batch.metrics)
|
|
|
|
batches.append(merged_batch)
|
|
return batches
|
|
|
|
def _log_run_data(self) -> None:
|
|
"""Process the run data in the running runs queues.
|
|
|
|
For each run in the running runs queues, this method retrieves the next batch of run data
|
|
from the queue and processes it by calling the `_processing_func` method with the run ID,
|
|
metrics, parameters, and tags in the batch. If the batch is empty, it is skipped. After
|
|
processing the batch, the processed watermark is updated and the batch event is set.
|
|
If an exception occurs during processing, the exception is logged and the batch event is set
|
|
with the exception. If the queue is empty, it is ignored.
|
|
|
|
Returns: None
|
|
"""
|
|
async_logging_buffer_seconds = MLFLOW_ASYNC_LOGGING_BUFFERING_SECONDS.get()
|
|
try:
|
|
if async_logging_buffer_seconds:
|
|
self._stop_data_logging_thread_event.wait(async_logging_buffer_seconds)
|
|
run_batches = self._fetch_batch_from_queue()
|
|
else:
|
|
run_batches = [self._queue.get(timeout=1)]
|
|
except Empty:
|
|
# Ignore empty queue exception
|
|
return
|
|
|
|
def logging_func(run_batch):
|
|
try:
|
|
self._logging_func(
|
|
run_id=run_batch.run_id,
|
|
metrics=run_batch.metrics,
|
|
params=run_batch.params,
|
|
tags=run_batch.tags,
|
|
)
|
|
except Exception as e:
|
|
_logger.error(f"Run Id {run_batch.run_id}: Failed to log run data: Exception: {e}")
|
|
run_batch.exception = e
|
|
finally:
|
|
run_batch.complete()
|
|
|
|
for run_batch in run_batches:
|
|
try:
|
|
self._batch_logging_worker_threadpool.submit(logging_func, run_batch)
|
|
except Exception as e:
|
|
_logger.error(
|
|
f"Failed to submit batch for logging: {e}. Usually this means you are not "
|
|
"shutting down MLflow properly before exiting. Please make sure you are using "
|
|
"context manager, e.g., `with mlflow.start_run():` or call `mlflow.end_run()`"
|
|
"explicitly to terminate MLflow logging before exiting."
|
|
)
|
|
run_batch.exception = e
|
|
run_batch.complete()
|
|
|
|
def _wait_for_batch(self, batch: RunBatch) -> None:
|
|
"""Wait for the given batch to be processed by the logging thread.
|
|
|
|
Args:
|
|
batch: The batch to wait for.
|
|
|
|
Raises:
|
|
Exception: If an exception occurred while processing the batch.
|
|
"""
|
|
batch.completion_event.wait()
|
|
if batch.exception:
|
|
raise batch.exception
|
|
|
|
def __getstate__(self):
|
|
"""Return the state of the object for pickling.
|
|
|
|
This method is called by the `pickle` module when the object is being pickled. It returns a
|
|
dictionary containing the object's state, with non-picklable attributes removed.
|
|
|
|
Returns:
|
|
dict: A dictionary containing the object's state.
|
|
"""
|
|
state = self.__dict__.copy()
|
|
del state["_queue"]
|
|
del state["_lock"]
|
|
del state["_status"]
|
|
|
|
if "_run_data_logging_thread" in state:
|
|
del state["_run_data_logging_thread"]
|
|
if "_stop_data_logging_thread_event" in state:
|
|
del state["_stop_data_logging_thread_event"]
|
|
if "_batch_logging_thread" in state:
|
|
del state["_batch_logging_thread"]
|
|
if "_batch_logging_worker_threadpool" in state:
|
|
del state["_batch_logging_worker_threadpool"]
|
|
if "_batch_status_check_threadpool" in state:
|
|
del state["_batch_status_check_threadpool"]
|
|
|
|
return state
|
|
|
|
def __setstate__(self, state):
|
|
"""Set the state of the object from a given state dictionary.
|
|
|
|
It pops back the removed non-picklable attributes from `self.__getstate__()`.
|
|
|
|
Args:
|
|
state (dict): A dictionary containing the state of the object.
|
|
|
|
Returns:
|
|
None
|
|
"""
|
|
self.__dict__.update(state)
|
|
self._queue = Queue()
|
|
self._lock = threading.RLock()
|
|
self._status = QueueStatus.IDLE
|
|
self._batch_logging_thread = None
|
|
self._batch_logging_worker_threadpool = None
|
|
self._batch_status_check_threadpool = None
|
|
self._stop_data_logging_thread_event = threading.Event()
|
|
|
|
def log_batch_async(
|
|
self, run_id: str, params: list[Param], tags: list[RunTag], metrics: list[Metric]
|
|
) -> RunOperations:
|
|
"""Asynchronously logs a batch of run data (parameters, tags, and metrics).
|
|
|
|
Args:
|
|
run_id (str): The ID of the run to log data for.
|
|
params (list[mlflow.entities.Param]): A list of parameters to log for the run.
|
|
tags (list[mlflow.entities.RunTag]): A list of tags to log for the run.
|
|
metrics (list[mlflow.entities.Metric]): A list of metrics to log for the run.
|
|
|
|
Returns:
|
|
mlflow.utils.async_utils.RunOperations: An object that encapsulates the
|
|
asynchronous operation of logging the batch of run data.
|
|
The object contains a list of `concurrent.futures.Future` objects that can be used
|
|
to check the status of the operation and retrieve any exceptions
|
|
that occurred during the operation.
|
|
"""
|
|
from mlflow import MlflowException
|
|
|
|
if not self.is_active():
|
|
raise MlflowException("AsyncLoggingQueue is not activated.")
|
|
batch = RunBatch(
|
|
run_id=run_id,
|
|
params=params,
|
|
tags=tags,
|
|
metrics=metrics,
|
|
completion_event=threading.Event(),
|
|
)
|
|
self._queue.put(batch)
|
|
operation_future = self._batch_status_check_threadpool.submit(self._wait_for_batch, batch)
|
|
return RunOperations(operation_futures=[operation_future])
|
|
|
|
def is_active(self) -> bool:
|
|
return self._status == QueueStatus.ACTIVE
|
|
|
|
def is_idle(self) -> bool:
|
|
return self._status == QueueStatus.IDLE
|
|
|
|
def _set_up_logging_thread(self) -> None:
|
|
"""
|
|
Sets up the logging thread.
|
|
|
|
This method shouldn't be called directly without shutting down the async
|
|
logging first if an existing async logging exists, otherwise it might
|
|
hang the program.
|
|
"""
|
|
with self._lock:
|
|
self._batch_logging_thread = threading.Thread(
|
|
target=self._logging_loop,
|
|
name="MLflowAsyncLoggingLoop",
|
|
daemon=True,
|
|
)
|
|
self._batch_logging_worker_threadpool = ThreadPoolExecutor(
|
|
max_workers=MLFLOW_ASYNC_LOGGING_THREADPOOL_SIZE.get() or 10,
|
|
thread_name_prefix=ASYNC_LOGGING_WORKER_THREAD_PREFIX,
|
|
)
|
|
|
|
self._batch_status_check_threadpool = ThreadPoolExecutor(
|
|
max_workers=MLFLOW_ASYNC_LOGGING_THREADPOOL_SIZE.get() or 10,
|
|
thread_name_prefix=ASYNC_LOGGING_STATUS_CHECK_THREAD_PREFIX,
|
|
)
|
|
|
|
self._batch_logging_thread.start()
|
|
|
|
def activate(self) -> None:
|
|
"""Activates the async logging queue
|
|
|
|
1. Initializes queue draining thread.
|
|
2. Initializes threads for checking the status of logged batch.
|
|
3. Registering an atexit callback to ensure that any remaining log data
|
|
is flushed before the program exits.
|
|
|
|
If the queue is already activated, this method does nothing.
|
|
"""
|
|
with self._lock:
|
|
if self.is_active():
|
|
return
|
|
|
|
self._set_up_logging_thread()
|
|
atexit.register(self._at_exit_callback)
|
|
|
|
self._status = QueueStatus.ACTIVE
|