""" Defines an AsyncArtifactsLoggingQueue that provides async fashion artifact writes using queue based approach. """ import atexit import logging import threading from concurrent.futures import ThreadPoolExecutor from queue import Empty, Queue from typing import TYPE_CHECKING, Callable, Union from mlflow.utils.async_logging.run_artifact import RunArtifact from mlflow.utils.async_logging.run_operations import RunOperations if TYPE_CHECKING: import PIL.Image _logger = logging.getLogger(__name__) class AsyncArtifactsLoggingQueue: """ This is a queue based run data processor that queue incoming data and process it using a single worker thread. This class is used to process artifacts saving in async fashion. Args: logging_func: A callable function that takes in three arguments: - filename: The name of the artifact file. - artifact_path: The path to the artifact. - artifact: The artifact to be logged. """ def __init__( self, artifact_logging_func: Callable[[str, str, Union["PIL.Image.Image"]], None] ) -> None: self._queue: Queue[RunArtifact] = Queue() self._lock = threading.RLock() self._artifact_logging_func = artifact_logging_func self._stop_data_logging_thread_event = threading.Event() self._is_activated = False 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 artifact processing status check. """ try: # Stop the data processing thread self._stop_data_logging_thread_event.set() # Waits till logging queue is drained. self._artifact_logging_thread.join() self._artifact_logging_worker_threadpool.shutdown(wait=True) self._artifact_status_check_threadpool.shutdown(wait=True) except Exception as e: _logger.error(f"Encountered error while trying to finish logging: {e}") def flush(self) -> None: """Flush the async logging queue. Calling this method will flush the queue to ensure all the data are logged. """ # Stop the data processing thread. self._stop_data_logging_thread_event.set() # Waits till logging queue is drained. self._artifact_logging_thread.join() self._artifact_logging_worker_threadpool.shutdown(wait=True) self._artifact_status_check_threadpool.shutdown(wait=True) # Restart the thread to listen to incoming data after flushing. self._stop_data_logging_thread_event.clear() self._set_up_logging_thread() 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_artifact() # Drain the queue after the stop event is set. while not self._queue.empty(): self._log_artifact() except Exception as e: from mlflow.exceptions import MlflowException raise MlflowException(f"Exception inside the run data logging thread: {e}") def _log_artifact(self) -> None: """Process the run's artifacts in the running runs queues. For each run in the running runs queues, this method retrieves the next artifact of run from the queue and processes it by calling the `_artifact_logging_func` method with the run ID and artifact. If the artifact is empty, it is skipped. After processing the artifact, the processed watermark is updated and the artifact event is set. If an exception occurs during processing, the exception is logged and the artifact event is set with the exception. If the queue is empty, it is ignored. """ try: run_artifact = self._queue.get(timeout=1) except Empty: # Ignore empty queue exception return def logging_func(run_artifact): try: self._artifact_logging_func( filename=run_artifact.filename, artifact_path=run_artifact.artifact_path, artifact=run_artifact.artifact, ) # Signal the artifact processing is done. run_artifact.completion_event.set() except Exception as e: _logger.error(f"Failed to log artifact {run_artifact.filename}. Exception: {e}") run_artifact.exception = e run_artifact.completion_event.set() self._artifact_logging_worker_threadpool.submit(logging_func, run_artifact) def _wait_for_artifact(self, artifact: RunArtifact) -> None: """Wait for given artifacts to be processed by the logging thread. Args: artifact: The artifact to wait for. Raises: Exception: If an exception occurred while processing the artifact. """ artifact.completion_event.wait() if artifact.exception: raise artifact.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["_is_activated"] if "_stop_data_logging_thread_event" in state: del state["_stop_data_logging_thread_event"] if "_artifact_logging_thread" in state: del state["_artifact_logging_thread"] if "_artifact_logging_worker_threadpool" in state: del state["_artifact_logging_worker_threadpool"] if "_artifact_status_check_threadpool" in state: del state["_artifact_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._is_activated = False self._artifact_logging_thread = None self._artifact_logging_worker_threadpool = None self._artifact_status_check_threadpool = None self._stop_data_logging_thread_event = threading.Event() def log_artifacts_async(self, filename, artifact_path, artifact) -> RunOperations: """Asynchronously logs runs artifacts. Args: filename: Filename of the artifact to be logged. artifact_path: Directory within the run's artifact directory in which to log the artifact. artifact: The artifact to be logged. Returns: mlflow.utils.async_utils.RunOperations: An object that encapsulates the asynchronous operation of logging the artifact 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_activated: raise MlflowException("AsyncArtifactsLoggingQueue is not activated.") artifact = RunArtifact( filename=filename, artifact_path=artifact_path, artifact=artifact, completion_event=threading.Event(), ) self._queue.put(artifact) operation_future = self._artifact_status_check_threadpool.submit( self._wait_for_artifact, artifact ) return RunOperations(operation_futures=[operation_future]) def is_active(self) -> bool: return self._is_activated def _set_up_logging_thread(self) -> None: """Sets up the logging thread. If the logging thread is already set up, this method does nothing. """ with self._lock: self._artifact_logging_thread = threading.Thread( target=self._logging_loop, name="MLflowAsyncArtifactsLoggingLoop", daemon=True, ) self._artifact_logging_worker_threadpool = ThreadPoolExecutor( max_workers=5, thread_name_prefix="MLflowArtifactsLoggingWorkerPool", ) self._artifact_status_check_threadpool = ThreadPoolExecutor( max_workers=5, thread_name_prefix="MLflowAsyncArtifactsLoggingStatusCheck", ) self._artifact_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 artifact. 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_activated: return self._set_up_logging_thread() atexit.register(self._at_exit_callback) self._is_activated = True