genesis-flow 1.0.0__py3-none-any.whl

mlflow/utils/autologging_utils/events.py

@@ -0,0 +1,294 @@
+import warnings
+from typing import Any
+
+from mlflow.utils.autologging_utils import _logger
+
+
+def _catch_exception(fn):
+    """A decorator that catches exceptions thrown by the wrapped function and logs them."""
+
+    def wrapper(*args):
+        try:
+            fn(*args)
+        except Exception as e:
+            _logger.debug(f"Failed to log autologging event via '{fn}'. Exception: {e}")
+
+    return wrapper
+
+
+class AutologgingEventLoggerWrapper:
+    """
+    A wrapper around AutologgingEventLogger for DRY:
+      - Store common arguments to avoid passing them to each logger method
+      - Catches exceptions thrown by the logger and logs them
+
+    NB: We could not modify the AutologgingEventLogger class directly because
+        it is used in Databricks code base as well.
+    """
+
+    def __init__(self, session, destination: Any, function_name: str):
+        self._session = session
+        self._destination = destination
+        self._function_name = function_name
+        self._logger = AutologgingEventLogger.get_logger()
+
+    @_catch_exception
+    def log_patch_function_start(self, args, kwargs):
+        self._logger.log_patch_function_start(
+            self._session, self._destination, self._function_name, args, kwargs
+        )
+
+    @_catch_exception
+    def log_patch_function_success(self, args, kwargs):
+        self._logger.log_patch_function_success(
+            self._session, self._destination, self._function_name, args, kwargs
+        )
+
+    @_catch_exception
+    def log_patch_function_error(self, args, kwargs, exception):
+        self._logger.log_patch_function_error(
+            self._session, self._destination, self._function_name, args, kwargs, exception
+        )
+
+    @_catch_exception
+    def log_original_function_start(self, args, kwargs):
+        self._logger.log_original_function_start(
+            self._session, self._destination, self._function_name, args, kwargs
+        )
+
+    @_catch_exception
+    def log_original_function_success(self, args, kwargs):
+        self._logger.log_original_function_success(
+            self._session, self._destination, self._function_name, args, kwargs
+        )
+
+    @_catch_exception
+    def log_original_function_error(self, args, kwargs, exception):
+        self._logger.log_original_function_error(
+            self._session, self._destination, self._function_name, args, kwargs, exception
+        )
+
+
+class AutologgingEventLogger:
+    """
+    Provides instrumentation hooks for important autologging lifecycle events, including:
+
+        - Calls to `mlflow.autolog()` APIs
+        - Calls to patched APIs with associated termination states
+          ("success" and "failure due to error")
+        - Calls to original / underlying APIs made by patched function code with
+          associated termination states ("success" and "failure due to error")
+
+    Default implementations are included for each of these hooks, which emit corresponding
+    DEBUG-level logging statements. Developers can provide their own hook implementations
+    by subclassing `AutologgingEventLogger` and calling the static
+    `AutologgingEventLogger.set_logger()` method to supply a new event logger instance.
+
+    Callers fetch the configured logger via `AutologgingEventLogger.get_logger()`
+    and invoke one or more hooks (e.g., `AutologgingEventLogger.get_logger().log_autolog_called()`).
+    """
+
+    _event_logger = None
+
+    @staticmethod
+    def get_logger():
+        """Fetches the configured `AutologgingEventLogger` instance for logging.
+
+        Returns:
+            The instance of `AutologgingEventLogger` specified via `set_logger`
+            (if configured) or the default implementation of `AutologgingEventLogger`
+            (if a logger was not configured via `set_logger`).
+
+        """
+        return AutologgingEventLogger._event_logger or AutologgingEventLogger()
+
+    @staticmethod
+    def set_logger(logger):
+        """Configures the `AutologgingEventLogger` instance for logging. This instance
+        is exposed via `AutologgingEventLogger.get_logger()` and callers use it to invoke
+        logging hooks (e.g., AutologgingEventLogger.get_logger().log_autolog_called()).
+
+        Args:
+            logger: The instance of `AutologgingEventLogger` to use when invoking logging hooks.
+
+        """
+        AutologgingEventLogger._event_logger = logger
+
+    def log_autolog_called(self, integration, call_args, call_kwargs):
+        """Called when the `autolog()` method for an autologging integration
+        is invoked (e.g., when a user invokes `mlflow.sklearn.autolog()`)
+
+        Args:
+            integration: The autologging integration for which `autolog()` was called.
+            call_args: **DEPRECATED** The positional arguments passed to the `autolog()` call.
+                This field is empty in MLflow > 1.13.1; all arguments are passed in
+                keyword form via `call_kwargs`.
+            call_kwargs: The arguments passed to the `autolog()` call in keyword form.
+                Any positional arguments should also be converted to keyword form
+                and passed via `call_kwargs`.
+        """
+        if len(call_args) > 0:
+            warnings.warn(
+                f"Received {len(call_args)} positional arguments via `call_args`. `call_args` is"
+                " deprecated in MLflow > 1.13.1, and all arguments should be passed"
+                " in keyword form via `call_kwargs`.",
+                category=DeprecationWarning,
+                stacklevel=2,
+            )
+        _logger.debug(
+            "Called autolog() method for %s autologging with args '%s' and kwargs '%s'",
+            integration,
+            call_args,
+            call_kwargs,
+        )
+
+    def log_patch_function_start(self, session, patch_obj, function_name, call_args, call_kwargs):
+        """Called upon invocation of a patched API associated with an autologging integration
+        (e.g., `sklearn.linear_model.LogisticRegression.fit()`).
+
+        Args:
+            session: The `AutologgingSession` associated with the patched API call.
+            patch_obj: The object (class, module, etc) on which the patched API was called.
+            function_name: The name of the patched API that was called.
+            call_args: The positional arguments passed to the patched API call.
+            call_kwargs: The keyword arguments passed to the patched API call.
+
+        """
+        _logger.debug(
+            "Invoked patched API '%s.%s' for %s autologging with args '%s' and kwargs '%s'",
+            patch_obj,
+            function_name,
+            session.integration,
+            call_args,
+            call_kwargs,
+        )
+
+    def log_patch_function_success(self, session, patch_obj, function_name, call_args, call_kwargs):
+        """
+        Called upon successful termination of a patched API associated with an autologging
+        integration (e.g., `sklearn.linear_model.LogisticRegression.fit()`).
+
+        Args:
+            session: The `AutologgingSession` associated with the patched API call.
+            patch_obj: The object (class, module, etc) on which the patched API was called.
+            function_name: The name of the patched API that was called.
+            call_args: The positional arguments passed to the patched API call.
+            call_kwargs: The keyword arguments passed to the patched API call.
+        """
+        _logger.debug(
+            "Patched API call '%s.%s' for %s autologging completed successfully. Patched ML"
+            " API was called with args '%s' and kwargs '%s'",
+            patch_obj,
+            function_name,
+            session.integration,
+            call_args,
+            call_kwargs,
+        )
+
+    def log_patch_function_error(
+        self, session, patch_obj, function_name, call_args, call_kwargs, exception
+    ):
+        """Called when execution of a patched API associated with an autologging integration
+        (e.g., `sklearn.linear_model.LogisticRegression.fit()`) terminates with an exception.
+
+        Args:
+            session: The `AutologgingSession` associated with the patched API call.
+            patch_obj: The object (class, module, etc) on which the patched API was called.
+            function_name: The name of the patched API that was called.
+            call_args: The positional arguments passed to the patched API call.
+            call_kwargs: The keyword arguments passed to the patched API call.
+            exception: The exception that caused the patched API call to terminate.
+        """
+        _logger.debug(
+            "Patched API call '%s.%s' for %s autologging threw exception. Patched API was"
+            " called with args '%s' and kwargs '%s'. Exception: %s",
+            patch_obj,
+            function_name,
+            session.integration,
+            call_args,
+            call_kwargs,
+            exception,
+        )
+
+    def log_original_function_start(
+        self, session, patch_obj, function_name, call_args, call_kwargs
+    ):
+        """
+        Called during the execution of a patched API associated with an autologging integration
+        when the original / underlying API is invoked. For example, this is called when
+        a patched implementation of `sklearn.linear_model.LogisticRegression.fit()` invokes
+        the original implementation of `sklearn.linear_model.LogisticRegression.fit()`.
+
+        Args:
+            session: The `AutologgingSession` associated with the patched API call.
+            patch_obj: The object (class, module, etc) on which the original API was called.
+            function_name: The name of the original API that was called.
+            call_args: The positional arguments passed to the original API call.
+            call_kwargs: The keyword arguments passed to the original API call.
+        """
+        _logger.debug(
+            "Original function invoked during execution of patched API '%s.%s' for %s"
+            " autologging. Original function was invoked with args '%s' and kwargs '%s'",
+            patch_obj,
+            function_name,
+            session.integration,
+            call_args,
+            call_kwargs,
+        )
+
+    def log_original_function_success(
+        self, session, patch_obj, function_name, call_args, call_kwargs
+    ):
+        """Called during the execution of a patched API associated with an autologging integration
+        when the original / underlying API invocation terminates successfully. For example,
+        when a patched implementation of `sklearn.linear_model.LogisticRegression.fit()` invokes the
+        original / underlying implementation of `LogisticRegression.fit()`, then this function is
+        called if the original / underlying implementation successfully completes.
+
+        Args:
+            session: The `AutologgingSession` associated with the patched API call.
+            patch_obj: The object (class, module, etc) on which the original API was called.
+            function_name: The name of the original API that was called.
+            call_args: The positional arguments passed to the original API call.
+            call_kwargs: The keyword arguments passed to the original API call.
+
+        """
+        _logger.debug(
+            "Original function invocation completed successfully during execution of patched API"
+            " call '%s.%s' for %s autologging. Original function was invoked with with"
+            " args '%s' and kwargs '%s'",
+            patch_obj,
+            function_name,
+            session.integration,
+            call_args,
+            call_kwargs,
+        )
+
+    def log_original_function_error(
+        self, session, patch_obj, function_name, call_args, call_kwargs, exception
+    ):
+        """Called during the execution of a patched API associated with an autologging integration
+        when the original / underlying API invocation terminates with an error. For example,
+        when a patched implementation of `sklearn.linear_model.LogisticRegression.fit()` invokes the
+        original / underlying implementation of `LogisticRegression.fit()`, then this function is
+        called if the original / underlying implementation terminates with an exception.
+
+        Args:
+            session: The `AutologgingSession` associated with the patched API call.
+            patch_obj: The object (class, module, etc) on which the original API was called.
+            function_name: The name of the original API that was called.
+            call_args: The positional arguments passed to the original API call.
+            call_kwargs: The keyword arguments passed to the original API call.
+            exception: The exception that caused the original API call to terminate.
+        """
+        _logger.debug(
+            "Original function invocation threw exception during execution of patched"
+            " API call '%s.%s' for %s autologging. Original function was invoked with"
+            " args '%s' and kwargs '%s'. Exception: %s",
+            patch_obj,
+            function_name,
+            session.integration,
+            call_args,
+            call_kwargs,
+            exception,
+        )

mlflow/utils/autologging_utils/logging_and_warnings.py

@@ -0,0 +1,328 @@
+import os
+import warnings
+from pathlib import Path
+from threading import RLock
+from threading import get_ident as get_current_thread_id
+
+import mlflow
+from mlflow.utils import logging_utils
+
+ORIGINAL_SHOWWARNING = warnings.showwarning
+
+
+class _WarningsController:
+    """
+    Provides threadsafe utilities to modify warning behavior for MLflow autologging, including:
+
+    - Global disablement of MLflow warnings across all threads
+    - Global rerouting of MLflow warnings to an MLflow event logger (i.e. `logger.warning()`)
+      across all threads
+    - Disablement of non-MLflow warnings for the current thread
+    - Rerouting of non-MLflow warnings to an MLflow event logger for the current thread
+    """
+
+    def __init__(self):
+        self._mlflow_root_path = Path(os.path.dirname(mlflow.__file__)).resolve()
+        self._state_lock = RLock()
+
+        self._did_patch_showwarning = False
+
+        self._disabled_threads = set()
+        self._rerouted_threads = set()
+        self._mlflow_warnings_disabled_globally = False
+        self._mlflow_warnings_rerouted_to_event_logs = False
+
+    def _patched_showwarning(self, message, category, filename, lineno, *args, **kwargs):
+        """
+        A patched implementation of `warnings.showwarning` that enforces the warning configuration
+        options configured on the controller (e.g. rerouting or disablement of MLflow warnings,
+        disablement of all warnings for the current thread).
+
+        Note that reassigning `warnings.showwarning` is the standard / recommended approach for
+        modifying warning message display behaviors. For reference, see
+        https://docs.python.org/3/library/warnings.html#warnings.showwarning
+        """
+        # NB: We explicitly avoid blocking on the `self._state_lock` lock during `showwarning`
+        # to so that threads don't have to execute serially whenever they emit warnings with
+        # `warnings.warn()`. We only lock during configuration changes to ensure that
+        # `warnings.showwarning` is patched or unpatched at the correct times.
+
+        from mlflow.utils.autologging_utils import _logger
+
+        # If the warning's source file is contained within the MLflow package's base
+        # directory, it is an MLflow warning and should be emitted via `logger.warning`
+        warning_source_path = Path(filename).resolve()
+        is_mlflow_warning = self._mlflow_root_path in warning_source_path.parents
+        curr_thread = get_current_thread_id()
+
+        if (curr_thread in self._disabled_threads) or (
+            is_mlflow_warning and self._mlflow_warnings_disabled_globally
+        ):
+            return
+        elif (curr_thread in self._rerouted_threads and not is_mlflow_warning) or (
+            is_mlflow_warning and self._mlflow_warnings_rerouted_to_event_logs
+        ):
+            _logger.warning(
+                'MLflow autologging encountered a warning: "%s:%d: %s: %s"',
+                filename,
+                lineno,
+                category.__name__,
+                message,
+            )
+        else:
+            ORIGINAL_SHOWWARNING(message, category, filename, lineno, *args, **kwargs)
+
+    def _should_patch_showwarning(self):
+        return (
+            (len(self._disabled_threads) > 0)
+            or (len(self._rerouted_threads) > 0)
+            or self._mlflow_warnings_disabled_globally
+            or self._mlflow_warnings_rerouted_to_event_logs
+        )
+
+    def _modify_patch_state_if_necessary(self):
+        """
+        Patches or unpatches `warnings.showwarning` if necessary, as determined by:
+            - Whether or not `warnings.showwarning` is already patched
+            - Whether or not any custom warning state has been configured on the warnings
+              controller (i.e. disablement or rerouting of certain warnings globally or for a
+              particular thread)
+
+        Note that reassigning `warnings.showwarning` is the standard / recommended approach for
+        modifying warning message display behaviors. For reference, see
+        https://docs.python.org/3/library/warnings.html#warnings.showwarning
+        """
+        with self._state_lock:
+            if self._should_patch_showwarning() and not self._did_patch_showwarning:
+                # NB: guard to prevent patching an instance of a patch
+                if warnings.showwarning != self._patched_showwarning:
+                    warnings.showwarning = self._patched_showwarning
+                self._did_patch_showwarning = True
+            elif not self._should_patch_showwarning() and self._did_patch_showwarning:
+                # NB: only unpatch iff the patched function is active
+                if warnings.showwarning == self._patched_showwarning:
+                    warnings.showwarning = ORIGINAL_SHOWWARNING
+                self._did_patch_showwarning = False
+
+    def set_mlflow_warnings_disablement_state_globally(self, disabled=True):
+        """Disables (or re-enables) MLflow warnings globally across all threads.
+
+        Args:
+            disabled: If `True`, disables MLflow warnings globally across all threads.
+                If `False`, enables MLflow warnings globally across all threads.
+
+        """
+        with self._state_lock:
+            self._mlflow_warnings_disabled_globally = disabled
+            self._modify_patch_state_if_necessary()
+
+    def set_mlflow_warnings_rerouting_state_globally(self, rerouted=True):
+        """
+        Enables (or disables) rerouting of MLflow warnings to an MLflow event logger with level
+        WARNING (e.g. `logger.warning()`) globally across all threads.
+
+        Args:
+            rerouted: If `True`, enables MLflow warning rerouting globally across all threads.
+                      If `False`, disables MLflow warning rerouting globally across all threads.
+        """
+        with self._state_lock:
+            self._mlflow_warnings_rerouted_to_event_logs = rerouted
+            self._modify_patch_state_if_necessary()
+
+    def set_non_mlflow_warnings_disablement_state_for_current_thread(self, disabled=True):
+        """Disables (or re-enables) non-MLflow warnings for the current thread.
+
+        Args:
+            disabled: If `True`, disables non-MLflow warnings for the current thread. If `False`,
+                enables non-MLflow warnings for the current thread. non-MLflow warning
+                behavior in other threads is unaffected.
+
+        """
+        with self._state_lock:
+            if disabled:
+                self._disabled_threads.add(get_current_thread_id())
+            else:
+                self._disabled_threads.discard(get_current_thread_id())
+            self._modify_patch_state_if_necessary()
+
+    def set_non_mlflow_warnings_rerouting_state_for_current_thread(self, rerouted=True):
+        """Enables (or disables) rerouting of non-MLflow warnings to an MLflow event logger with
+        level WARNING (e.g. `logger.warning()`) for the current thread.
+
+        Args:
+            rerouted: If `True`, enables non-MLflow warning rerouting for the current thread.
+                If `False`, disables non-MLflow warning rerouting for the current thread.
+                non-MLflow warning behavior in other threads is unaffected.
+
+        """
+        with self._state_lock:
+            if rerouted:
+                self._rerouted_threads.add(get_current_thread_id())
+            else:
+                self._rerouted_threads.discard(get_current_thread_id())
+            self._modify_patch_state_if_necessary()
+
+    def get_warnings_disablement_state_for_current_thread(self):
+        """
+        Returns:
+            True if non-MLflow warnings are disabled for the current thread. False otherwise.
+        """
+        return get_current_thread_id() in self._disabled_threads
+
+    def get_warnings_rerouting_state_for_current_thread(self):
+        """
+        Returns:
+            True if non-MLflow warnings are rerouted to an MLflow event logger with level
+            WARNING for the current thread. False otherwise.
+        """
+        return get_current_thread_id() in self._rerouted_threads
+
+
+_WARNINGS_CONTROLLER = _WarningsController()
+
+
+class NonMlflowWarningsBehaviorForCurrentThread:
+    """
+    Context manager that modifies the behavior of non-MLflow warnings upon entry, according to the
+    specified parameters.
+
+    Args:
+        disable_warnings: If `True`, disable  (mutate & discard) non-MLflow warnings. If `False`,
+            do not disable non-MLflow warnings.
+        reroute_warnings: If `True`, reroute non-MLflow warnings to an MLflow event logger with
+            level WARNING. If `False`, do not reroute non-MLflow warnings.
+    """
+
+    def __init__(self, disable_warnings, reroute_warnings):
+        self._disable_warnings = disable_warnings
+        self._reroute_warnings = reroute_warnings
+        self._prev_disablement_state = None
+        self._prev_rerouting_state = None
+
+    def __enter__(self):
+        self._enter_impl()
+        return self
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        self._exit_impl(exc_type, exc_val, exc_tb)
+
+    async def __aenter__(self):
+        self._enter_impl()
+        return self
+
+    async def __aexit__(self, exc_type, exc_val, exc_tb):
+        self._exit_impl(exc_type, exc_val, exc_tb)
+
+    def _enter_impl(self):
+        self._prev_disablement_state = (
+            _WARNINGS_CONTROLLER.get_warnings_disablement_state_for_current_thread()
+        )
+        self._prev_rerouting_state = (
+            _WARNINGS_CONTROLLER.get_warnings_rerouting_state_for_current_thread()
+        )
+        try:
+            _WARNINGS_CONTROLLER.set_non_mlflow_warnings_disablement_state_for_current_thread(
+                disabled=self._disable_warnings
+            )
+            _WARNINGS_CONTROLLER.set_non_mlflow_warnings_rerouting_state_for_current_thread(
+                rerouted=self._reroute_warnings
+            )
+        except Exception:
+            pass
+
+    def _exit_impl(self, *args, **kwargs):
+        _WARNINGS_CONTROLLER.set_non_mlflow_warnings_disablement_state_for_current_thread(
+            disabled=self._prev_disablement_state
+        )
+        _WARNINGS_CONTROLLER.set_non_mlflow_warnings_rerouting_state_for_current_thread(
+            rerouted=self._prev_rerouting_state
+        )
+
+
+class MlflowEventsAndWarningsBehaviorGlobally:
+    """
+    Threadsafe context manager that modifies the behavior of MLflow event logging statements
+    and MLflow warnings upon entry, according to the specified parameters. Modifications are
+    applied globally across all threads and are not reverted until all threads that have made
+    a particular modification have exited the context.
+
+    Args:
+        disable_event_logs: If `True`, disable (mute & discard) MLflow event logging statements.
+            If `False`, do not disable MLflow event logging statements.
+        disable_warnings: If `True`, disable  (mutate & discard) MLflow warnings. If `False`,
+            do not disable MLflow warnings.
+        reroute_warnings: If `True`, reroute MLflow warnings to an MLflow event logger with
+            level WARNING. If `False`, do not reroute MLflow warnings.
+
+    """
+
+    _lock = RLock()
+    _disable_event_logs_count = 0
+    _disable_warnings_count = 0
+    _reroute_warnings_count = 0
+
+    def __init__(self, disable_event_logs, disable_warnings, reroute_warnings):
+        self._disable_event_logs = disable_event_logs
+        self._disable_warnings = disable_warnings
+        self._reroute_warnings = reroute_warnings
+
+    def __enter__(self):
+        self._enter_impl()
+        return self
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        self._exit_impl(exc_type, exc_val, exc_tb)
+
+    async def __aenter__(self):
+        self._enter_impl()
+        return self
+
+    async def __aexit__(self, exc_type, exc_val, exc_tb):
+        self._exit_impl(exc_type, exc_val, exc_tb)
+
+    def _enter_impl(self):
+        try:
+            with MlflowEventsAndWarningsBehaviorGlobally._lock:
+                if self._disable_event_logs:
+                    if MlflowEventsAndWarningsBehaviorGlobally._disable_event_logs_count <= 0:
+                        logging_utils.disable_logging()
+                    MlflowEventsAndWarningsBehaviorGlobally._disable_event_logs_count += 1
+
+                if self._disable_warnings:
+                    if MlflowEventsAndWarningsBehaviorGlobally._disable_warnings_count <= 0:
+                        _WARNINGS_CONTROLLER.set_mlflow_warnings_disablement_state_globally(
+                            disabled=True
+                        )
+                    MlflowEventsAndWarningsBehaviorGlobally._disable_warnings_count += 1
+
+                if self._reroute_warnings:
+                    if MlflowEventsAndWarningsBehaviorGlobally._reroute_warnings_count <= 0:
+                        _WARNINGS_CONTROLLER.set_mlflow_warnings_rerouting_state_globally(
+                            rerouted=True
+                        )
+                    MlflowEventsAndWarningsBehaviorGlobally._reroute_warnings_count += 1
+        except Exception:
+            pass
+
+    def _exit_impl(self, *args, **kwargs):
+        try:
+            with MlflowEventsAndWarningsBehaviorGlobally._lock:
+                if self._disable_event_logs:
+                    MlflowEventsAndWarningsBehaviorGlobally._disable_event_logs_count -= 1
+                if self._disable_warnings:
+                    MlflowEventsAndWarningsBehaviorGlobally._disable_warnings_count -= 1
+                if self._reroute_warnings:
+                    MlflowEventsAndWarningsBehaviorGlobally._reroute_warnings_count -= 1
+
+                if MlflowEventsAndWarningsBehaviorGlobally._disable_event_logs_count <= 0:
+                    logging_utils.enable_logging()
+                if MlflowEventsAndWarningsBehaviorGlobally._disable_warnings_count <= 0:
+                    _WARNINGS_CONTROLLER.set_mlflow_warnings_disablement_state_globally(
+                        disabled=False
+                    )
+                if MlflowEventsAndWarningsBehaviorGlobally._reroute_warnings_count <= 0:
+                    _WARNINGS_CONTROLLER.set_mlflow_warnings_rerouting_state_globally(
+                        rerouted=False
+                    )
+        except Exception:
+            pass

mlflow/utils/autologging_utils/metrics_queue.py

@@ -0,0 +1,71 @@
+import concurrent.futures
+from threading import RLock
+
+from mlflow.entities import Metric
+from mlflow.tracking.client import MlflowClient
+
+_metrics_queue_lock = RLock()
+_metrics_queue = []
+_thread_pool = concurrent.futures.ThreadPoolExecutor(
+    max_workers=1, thread_name_prefix="MlflowMetricsQueue"
+)
+
+_MAX_METRIC_QUEUE_SIZE = 500
+
+
+def _assoc_list_to_map(lst):
+    """
+    Convert an association list to a dictionary.
+    """
+    d = {}
+    for run_id, metric in lst:
+        d[run_id] = d[run_id] + [metric] if run_id in d else [metric]
+    return d
+
+
+def flush_metrics_queue():
+    """Flush the metric queue and log contents in batches to MLflow.
+
+    Queue is divided into batches according to run id.
+    """
+    try:
+        # Multiple queue flushes may be scheduled simultaneously on different threads
+        # (e.g., if the queue is at its flush threshold and several more items
+        # are added before a flush occurs). For correctness and efficiency, only one such
+        # flush operation should proceed; all others are redundant and should be dropped
+        acquired_lock = _metrics_queue_lock.acquire(blocking=False)
+        if acquired_lock:
+            client = MlflowClient()
+            # For thread safety and to avoid modifying a list while iterating over it, we record a
+            # separate list of the items being flushed and remove each one from the metric queue,
+            # rather than clearing the metric queue or reassigning it (clearing / reassigning is
+            # dangerous because we don't block threads from adding to the queue while a flush is
+            # in progress)
+            snapshot = _metrics_queue[:]
+            for item in snapshot:
+                _metrics_queue.remove(item)
+
+            metrics_by_run = _assoc_list_to_map(snapshot)
+            for run_id, metrics in metrics_by_run.items():
+                client.log_batch(run_id, metrics=metrics, params=[], tags=[])
+    finally:
+        if acquired_lock:
+            _metrics_queue_lock.release()
+
+
+def add_to_metrics_queue(key, value, step, time, run_id):
+    """Add a metric to the metric queue.
+
+    Flush the queue if it exceeds max size.
+
+    Args:
+        key: string, the metrics key,
+        value: float, the metrics value.
+        step: int, the step of current metric.
+        time: int, the timestamp of current metric.
+        run_id: string, the run id of the associated mlflow run.
+    """
+    met = Metric(key=key, value=value, timestamp=time, step=step)
+    _metrics_queue.append((run_id, met))
+    if len(_metrics_queue) > _MAX_METRIC_QUEUE_SIZE:
+        _thread_pool.submit(flush_metrics_queue)