Add silent option for autologging integrations (mlflow#4173)

* Rename file Signed-off-by: dbczumar <corey.zumar@databricks.com> * Rename imports Signed-off-by: dbczumar <corey.zumar@databricks.com> * More Signed-off-by: dbczumar <corey.zumar@databricks.com> * More Signed-off-by: dbczumar <corey.zumar@databricks.com> * More cleanup Signed-off-by: dbczumar <corey.zumar@databricks.com> * Format Signed-off-by: dbczumar <corey.zumar@databricks.com> * Repath Signed-off-by: dbczumar <corey.zumar@databricks.com> * Partial checkout Signed-off-by: dbczumar <corey.zumar@databricks.com> * XGB and more Signed-off-by: dbczumar <corey.zumar@databricks.com> * Reformat Signed-off-by: dbczumar <corey.zumar@databricks.com> * Add needed utils Signed-off-by: dbczumar <corey.zumar@databricks.com> * Wrap patch Signed-off-by: dbczumar <corey.zumar@databricks.com> * Format Signed-off-by: dbczumar <corey.zumar@databricks.com> * More checkout Signed-off-by: dbczumar <corey.zumar@databricks.com> * Yet another format Signed-off-by: dbczumar <corey.zumar@databricks.com> * Add back some defaults Signed-off-by: dbczumar <corey.zumar@databricks.com> * Warning disable consolidate Signed-off-by: dbczumar <corey.zumar@databricks.com> * Fix some tests Signed-off-by: dbczumar <corey.zumar@databricks.com> * More docs Signed-off-by: dbczumar <corey.zumar@databricks.com> * Fixture fixes and fluent test Signed-off-by: dbczumar <corey.zumar@databricks.com> * A little more doc clarification Signed-off-by: dbczumar <corey.zumar@databricks.com> * Remove redundant test file, update conftest Signed-off-by: dbczumar <corey.zumar@databricks.com> * Clarifying comment Signed-off-by: dbczumar <corey.zumar@databricks.com> * Remove duplicate fixture Signed-off-by: dbczumar <corey.zumar@databricks.com> * Simplifications & test cases Signed-off-by: dbczumar <corey.zumar@databricks.com> * Docs Signed-off-by: dbczumar <corey.zumar@databricks.com> * Format Signed-off-by: dbczumar <corey.zumar@databricks.com> * Lint fixes Signed-off-by: dbczumar <corey.zumar@databricks.com> * Simplify and format Signed-off-by: dbczumar <corey.zumar@databricks.com> * Remove unused import Signed-off-by: dbczumar <corey.zumar@databricks.com> * Refs Signed-off-by: dbczumar <corey.zumar@databricks.com> * Format Signed-off-by: dbczumar <corey.zumar@databricks.com> * Remove TestStream Signed-off-by: dbczumar <corey.zumar@databricks.com> * Address remaining Jin comments Signed-off-by: dbczumar <corey.zumar@databricks.com> * Most of a test case Signed-off-by: dbczumar <corey.zumar@databricks.com> * Finish test case Signed-off-by: dbczumar <corey.zumar@databricks.com>
simonvanbernem · Mar 26, 2021 · a777160 · a777160
1 parent a4df316
commit a777160
Show file tree

Hide file tree

Showing 20 changed files with 1,052 additions and 320 deletions.
diff --git a/conftest.py b/conftest.py
@@ -93,8 +93,9 @@ def pytest_ignore_collect(path, config):
             # this test is included here because it imports many big libraries like tf, keras, etc
             "tests/tracking/fluent/test_fluent_autolog.py",
             # cross flavor autologging related tests.
-            "tests/autologging/test_autologging_safety_integration.py",
             "tests/autologging/test_autologging_safety_unit.py",
+            "tests/autologging/test_autologging_behaviors_unit.py",
+            "tests/autologging/test_autologging_behaviors_integration.py",
             "tests/autologging/test_autologging_utils.py",
         ]
 

diff --git a/mlflow/fastai.py b/mlflow/fastai.py
@@ -381,7 +381,11 @@ def load_model(model_uri):
 @experimental
 @autologging_integration(FLAVOR_NAME)
 def autolog(
-    log_models=True, disable=False, exclusive=False, disable_for_unsupported_versions=False
+    log_models=True,
+    disable=False,
+    exclusive=False,
+    disable_for_unsupported_versions=False,
+    silent=False,
 ):  # pylint: disable=unused-argument
     """
     Enable automatic logging from Fastai to MLflow.
@@ -403,6 +407,9 @@ def autolog(
     :param disable_for_unsupported_versions: If ``True``, disable autologging for versions of
                       fastai that have not been tested against this version of the MLflow client
                       or are incompatible.
+    :param silent: If ``True``, suppress all event logs and warnings from MLflow during Fastai
+                   autologging. If ``False``, show all events and warnings during Fastai
+                   autologging.
 
     .. code-block:: python
         :caption: Example

diff --git a/mlflow/gluon.py b/mlflow/gluon.py
@@ -332,7 +332,11 @@ def log_model(
 @experimental
 @autologging_integration(FLAVOR_NAME)
 def autolog(
-    log_models=True, disable=False, exclusive=False, disable_for_unsupported_versions=False
+    log_models=True,
+    disable=False,
+    exclusive=False,
+    disable_for_unsupported_versions=False,
+    silent=False,
 ):  # pylint: disable=unused-argument
     """
     Enables (or disables) and configures autologging from Gluon to MLflow.
@@ -350,6 +354,9 @@ def autolog(
     :param disable_for_unsupported_versions: If ``True``, disable autologging for versions of
                       gluon that have not been tested against this version of the MLflow client
                       or are incompatible.
+    :param silent: If ``True``, suppress all event logs and warnings from MLflow during MXNet Gluon
+                   autologging. If ``False``, show all events and warnings during MXNet Gluon
+                   autologging.
     """
 
     from mxnet.gluon.contrib.estimator import Estimator, EpochEnd, TrainBegin, TrainEnd

diff --git a/mlflow/keras.py b/mlflow/keras.py
@@ -572,7 +572,11 @@ def load_model(model_uri, **kwargs):
 @experimental
 @autologging_integration(FLAVOR_NAME)
 def autolog(
-    log_models=True, disable=False, exclusive=False, disable_for_unsupported_versions=False
+    log_models=True,
+    disable=False,
+    exclusive=False,
+    disable_for_unsupported_versions=False,
+    silent=False,
 ):  # pylint: disable=unused-argument
     # pylint: disable=E0611
     """
@@ -632,6 +636,9 @@ def autolog(
     :param disable_for_unsupported_versions: If ``True``, disable autologging for versions of
                       keras that have not been tested against this version of the MLflow client
                       or are incompatible.
+    :param silent: If ``True``, suppress all event logs and warnings from MLflow during Keras
+                   autologging. If ``False``, show all events and warnings during Keras
+                   autologging.
     """
     import keras
 

diff --git a/mlflow/lightgbm.py b/mlflow/lightgbm.py
@@ -291,6 +291,7 @@ def autolog(
     disable=False,
     exclusive=False,
     disable_for_unsupported_versions=False,
+    silent=False,
 ):  # pylint: disable=unused-argument
     """
     Enables (or disables) and configures autologging from LightGBM to MLflow. Logs the following:
@@ -329,6 +330,9 @@ def autolog(
     :param disable_for_unsupported_versions: If ``True``, disable autologging for versions of
                       lightgbm that have not been tested against this version of the MLflow client
                       or are incompatible.
+    :param silent: If ``True``, suppress all event logs and warnings from MLflow during LightGBM
+                   autologging. If ``False``, show all events and warnings during LightGBM
+                   autologging.
     """
     import lightgbm
     import numpy as np

diff --git a/mlflow/pytorch/__init__.py b/mlflow/pytorch/__init__.py
@@ -836,6 +836,7 @@ def autolog(
     disable=False,
     exclusive=False,
     disable_for_unsupported_versions=False,
+    silent=False,
 ):  # pylint: disable=unused-argument
     """
     Enables (or disables) and configures autologging from `PyTorch Lightning
@@ -869,6 +870,9 @@ def autolog(
     :param disable_for_unsupported_versions: If ``True``, disable autologging for versions of
                       pytorch and pytorch-lightning that have not been tested against this version
                       of the MLflow client or are incompatible.
+    :param silent: If ``True``, suppress all event logs and warnings from MLflow during PyTorch
+                   Lightning autologging. If ``False``, show all events and warnings during
+                   PyTorch Lightning autologging.
 
     .. code-block:: python
         :caption: Example

diff --git a/mlflow/sklearn/__init__.py b/mlflow/sklearn/__init__.py
@@ -536,6 +536,7 @@ def autolog(
     disable=False,
     exclusive=False,
     disable_for_unsupported_versions=False,
+    silent=False,
 ):  # pylint: disable=unused-argument
     """
     Enables (or disables) and configures autologging for scikit-learn estimators.
@@ -724,6 +725,9 @@ def fetch_logged_data(run_id):
     :param disable_for_unsupported_versions: If ``True``, disable autologging for versions of
                       scikit-learn that have not been tested against this version of the MLflow
                       client or are incompatible.
+    :param silent: If ``True``, suppress all event logs and warnings from MLflow during scikit-learn
+                   autologging. If ``False``, show all events and warnings during scikit-learn
+                   autologging.
     """
     import pandas as pd
     import sklearn

diff --git a/mlflow/spark.py b/mlflow/spark.py
@@ -682,7 +682,7 @@ def predict(self, pandas_df):
 
 @experimental
 @autologging_integration(FLAVOR_NAME)
-def autolog(disable=False):  # pylint: disable=unused-argument
+def autolog(disable=False, silent=False):  # pylint: disable=unused-argument
     """
     Enables (or disables) and configures logging of Spark datasource paths, versions
     (if applicable), and formats when they are read. This method is not threadsafe and assumes a
@@ -739,6 +739,9 @@ def autolog(disable=False):  # pylint: disable=unused-argument
 
     :param disable: If ``True``, disables the Spark datasource autologging integration.
                     If ``False``, enables the Spark datasource autologging integration.
+    :param silent: If ``True``, suppress all event logs and warnings from MLflow during Spark
+                   datasource autologging. If ``False``, show all events and warnings during Spark
+                   datasource autologging.
     """
     from mlflow.utils._spark_utils import _get_active_spark_session
     from mlflow._spark_autologging import _listen_for_spark_activity

diff --git a/mlflow/statsmodels.py b/mlflow/statsmodels.py
@@ -311,7 +311,11 @@ class AutologHelpers:
 @experimental
 @autologging_integration(FLAVOR_NAME)
 def autolog(
-    log_models=True, disable=False, exclusive=False, disable_for_unsupported_versions=False
+    log_models=True,
+    disable=False,
+    exclusive=False,
+    disable_for_unsupported_versions=False,
+    silent=False,
 ):  # pylint: disable=unused-argument
     """
     Enables (or disables) and configures automatic logging from statsmodels to MLflow.
@@ -332,6 +336,9 @@ def autolog(
     :param disable_for_unsupported_versions: If ``True``, disable autologging for versions of
                       statsmodels that have not been tested against this version of the MLflow
                       client or are incompatible.
+    :param silent: If ``True``, suppress all event logs and warnings from MLflow during statsmodels
+                   autologging. If ``False``, show all events and warnings during statsmodels
+                   autologging.
     """
     import statsmodels
 

diff --git a/mlflow/tensorflow.py b/mlflow/tensorflow.py
@@ -826,6 +826,7 @@ def autolog(
     disable=False,
     exclusive=False,
     disable_for_unsupported_versions=False,
+    silent=False,
 ):  # pylint: disable=unused-argument
     # pylint: disable=E0611
     """
@@ -889,6 +890,9 @@ def autolog(
     :param disable_for_unsupported_versions: If ``True``, disable autologging for versions of
                       tensorflow that have not been tested against this version of the MLflow
                       client or are incompatible.
+    :param silent: If ``True``, suppress all event logs and warnings from MLflow during TensorFlow
+                   autologging. If ``False``, show all events and warnings during TensorFlow
+                   autologging.
     """
     import tensorflow
 

diff --git a/mlflow/tracking/fluent.py b/mlflow/tracking/fluent.py
@@ -1227,6 +1227,7 @@ def autolog(
     disable=False,
     exclusive=False,
     disable_for_unsupported_versions=False,
+    silent=False,
 ):  # pylint: disable=unused-argument
     """
     Enables (or disables) and configures autologging for all supported integrations.
@@ -1282,6 +1283,9 @@ def autolog(
     :param disable_for_unsupported_versions: If ``True``, disable autologging for versions of
                       all integration libraries that have not been tested against this version
                       of the MLflow client or are incompatible.
+    :param silent: If ``True``, suppress all event logs and warnings from MLflow during autologging
+                   setup and training execution. If ``False``, show all events and warnings during
+                   autologging setup and training execution.
 
     .. code-block:: python
         :caption: Example
@@ -1395,14 +1399,16 @@ def setup_autologging(module):
             AUTOLOGGING_INTEGRATIONS[autolog_fn.integration_name][
                 CONF_KEY_IS_GLOBALLY_CONFIGURED
             ] = True
-            if not autologging_is_disabled(autolog_fn.integration_name):
+            if not autologging_is_disabled(
+                autolog_fn.integration_name
+            ) and not autologging_params.get("silent", False):
                 _logger.info("Autologging successfully enabled for %s.", module.__name__)
         except Exception as e:
             if _is_testing():
                 # Raise unexpected exceptions in test mode in order to detect
                 # errors within dependent autologging integrations
                 raise
-            else:
+            elif not autologging_params.get("silent", False):
                 _logger.warning(
                     "Exception raised while enabling autologging for %s: %s",
                     module.__name__,