TensorFlow autolog automatic run ending (mlflow#2094)

* added keras logic and tests

* refactored keras autolog tests

* minor edits to keras tests

* minor edit to documentation

* added autolog ending for tf.Keras, added patching for tf.keras.model.fit_generator

* refactored and added tests for tf.keras

* removed fit_generator hook for another PR

* renamed create_model to create_keras_model

* refactored and added tests for auto ending tf.keras autolog runs

* made test names more descriptive

* lint

* added run starting in keras

* refactored run starting into yield fixture

* added logic for run autoending in tf.estimator

* added logic and tests for tf estimator run ending

* added _autolog_run_id, actually patched train

* fixed wrapper tries, refactored tests

* forgot a couple of parentheses

* updated docs

* updated docs more, removed extra line in tests

* changed wording for docs

* added tests for tf1

* added space in docs

* moved space in docs

* removed code in docs

* grammar in docs

* added param checking test

* fixed malformed table

* removing _AUTO_END_RUN

* added check for ending auto runs

* replaced manual tempdir with tmpdir fixture

* removing export_savedmodel from docs

Co-Authored-By: Siddharth Murching <smurching@gmail.com>

* Update docs wording

Co-Authored-By: Siddharth Murching <smurching@gmail.com>

* Added defensive check to run ending

Co-Authored-By: Siddharth Murching <smurching@gmail.com>

* created contextmanager

* fixed defensive check

* lint

* edited docs

Author: juntai-zheng

docs/source/tracking.rst

@@ -217,20 +217,26 @@ log statements. See example usages with `Keras <https://github.com/mlflow/mlflow
 
 Autologging captures the following information:
 
-+------------------+--------------------------------------------------------+----------------------------------------------------------+---------------+-------------------------------------------------------------------------------------------------------------------------------+
-| Framework        | Metrics                                                | Parameters                                               | Tags          | Artifacts                                                                                                                     |
-+------------------+--------------------------------------------------------+----------------------------------------------------------+---------------+-------------------------------------------------------------------------------------------------------------------------------+
-| Keras            | Training loss; validation loss; user-specified metrics | Number of layers; optimizer name; learning rate; epsilon | Model summary | `MLflow Model <https://mlflow.org/docs/latest/models.html>`_ (Keras model); on training end                                   |
-+------------------+--------------------------------------------------------+----------------------------------------------------------+---------------+-------------------------------------------------------------------------------------------------------------------------------+
-| ``tf.keras``     | Training loss; validation loss; user-specified metrics | Number of layers; optimizer name; learning rate; epsilon | Model summary | `MLflow Model <https://mlflow.org/docs/latest/models.html>`_ (Keras model), TensorBoard logs; on training end                 |
-+------------------+--------------------------------------------------------+----------------------------------------------------------+---------------+-------------------------------------------------------------------------------------------------------------------------------+
-| ``tf.estimator`` | TensorBoard metrics                                    | --                                                       | --            | `MLflow Model <https://mlflow.org/docs/latest/models.html>`_ (TF saved model); on call to ``tf.estimator.export_saved_model`` |
-+------------------+--------------------------------------------------------+----------------------------------------------------------+---------------+-------------------------------------------------------------------------------------------------------------------------------+
-| TensorFlow Core  | All ``tf.summary.scalar`` calls                        | --                                                       | --            | --                                                                                                                            |
-+------------------+--------------------------------------------------------+----------------------------------------------------------+---------------+-------------------------------------------------------------------------------------------------------------------------------+
++------------------+--------------------------------------------------------+----------------------------------------------------------+---------------+--------------------------------------------------------------------------------------------------------------------------------------------------+
+| Framework        | Metrics                                                | Parameters                                               | Tags          | Artifacts                                                                                                                                        |
++------------------+--------------------------------------------------------+----------------------------------------------------------+---------------+--------------------------------------------------------------------------------------------------------------------------------------------------+
+| Keras            | Training loss; validation loss; user-specified metrics | Number of layers; optimizer name; learning rate; epsilon | Model summary | Model summary on training start; `MLflow Model <https://mlflow.org/docs/latest/models.html>`_ (Keras model) on training end                      |
++------------------+--------------------------------------------------------+----------------------------------------------------------+---------------+--------------------------------------------------------------------------------------------------------------------------------------------------+
+| ``tf.keras``     | Training loss; validation loss; user-specified metrics | Number of layers; optimizer name; learning rate; epsilon | Model summary | Model summary on training start; `MLflow Model <https://mlflow.org/docs/latest/models.html>`_ (Keras model), TensorBoard logs on training end    |
++------------------+--------------------------------------------------------+----------------------------------------------------------+---------------+--------------------------------------------------------------------------------------------------------------------------------------------------+
+| ``tf.estimator`` | TensorBoard metrics                                    | steps, max_steps                                         | --            | `MLflow Model <https://mlflow.org/docs/latest/models.html>`_ (TF saved model) on call to ``tf.estimator.export_saved_model``                     |
++------------------+--------------------------------------------------------+----------------------------------------------------------+---------------+--------------------------------------------------------------------------------------------------------------------------------------------------+
+| TensorFlow Core  | All ``tf.summary.scalar`` calls                        | --                                                       | --            | --                                                                                                                                               |
++------------------+--------------------------------------------------------+----------------------------------------------------------+---------------+--------------------------------------------------------------------------------------------------------------------------------------------------+
 
 Note that autologging for ``tf.keras`` is handled by :py:func:`mlflow.tensorflow.autolog`, not :py:func:`mlflow.keras.autolog`. 
 
+If no active run exists when ``autolog()`` captures data, MLflow will automatically create a run to log information to.
+Once training ends via calls to ``tf.estimator.train()``, ``tf.keras.fit()``, ``tf.keras.fit_generator()``, ``keras.fit()`` or ``keras.fit_generator()``,
+or once ``tf.estimator`` models are exported via ``tf.estimator.export_saved_model()``, MLflow will automatically end that run.
+
+If a run exists when ``autolog()`` captures data, MLflow will log to that run and not automatically end that run after training.
+
 **Note**: this feature is experimental - the API and format of the logged data are subject to change.
 
 

mlflow/tensorflow.py

@@ -27,6 +27,7 @@
 import tensorflow
 import mlflow.keras
 from distutils.version import LooseVersion
+from contextlib import contextmanager
 from tensorflow.keras.callbacks import Callback, TensorBoard  # pylint: disable=import-error
 from mlflow import pyfunc
 from mlflow.exceptions import MlflowException
@@ -54,7 +55,7 @@
 _thread_pool = concurrent.futures.ThreadPoolExecutor(max_workers=1)
 
 # For tracking if the run was started by autologging.
-_AUTO_END_RUN = False
+_AUTOLOG_RUN_ID = None
 
 
 def get_default_conda_env():
@@ -584,8 +585,8 @@ def _log_event(event):
     """
     if not mlflow.active_run():
         try_mlflow_log(mlflow.start_run)
-        global _AUTO_END_RUN
-        _AUTO_END_RUN = True
+        global _AUTOLOG_RUN_ID
+        _AUTOLOG_RUN_ID = mlflow.active_run().info.run_id
     if event.WhichOneof('what') == 'summary':
         summary = event.summary
         for v in summary.value:
@@ -657,85 +658,123 @@ def autolog(every_n_iter=100):
                       "1.12 <= v <= 2.0.0 are supported.")
         return
 
+    @contextmanager
+    def _manage_active_run():
+        if not mlflow.active_run():
+            try_mlflow_log(mlflow.start_run)
+            global _AUTOLOG_RUN_ID
+            if mlflow.active_run() is not None:  # defensive check in case `mlflow.start_run` fails
+                _AUTOLOG_RUN_ID = mlflow.active_run().info.run_id
+        yield mlflow.active_run()
+        if mlflow.active_run() is not None and mlflow.active_run().info.run_id == _AUTOLOG_RUN_ID:
+            try_mlflow_log(mlflow.end_run)
+
+    @gorilla.patch(tensorflow.estimator.Estimator)
+    def train(self, *args, **kwargs):
+        with _manage_active_run():
+            original = gorilla.get_original_attribute(tensorflow.estimator.Estimator, 'train')
+
+            # Checking step and max_step parameters for logging
+            if len(args) >= 3:
+                try_mlflow_log(mlflow.log_param, 'steps', args[2])
+                if len(args) >= 4:
+                    try_mlflow_log(mlflow.log_param, 'max_steps', args[3])
+            if 'steps' in kwargs:
+                try_mlflow_log(mlflow.log_param, 'steps', kwargs['steps'])
+            if 'max_steps' in kwargs:
+                try_mlflow_log(mlflow.log_param, 'max_steps', kwargs['max_steps'])
+
+            result = original(self, *args, **kwargs)
+
+            return result
+
     @gorilla.patch(tensorflow.estimator.Estimator)
     def export_saved_model(self, *args, **kwargs):
+        auto_end = False
+        if not mlflow.active_run():
+            global _AUTOLOG_RUN_ID
+            if _AUTOLOG_RUN_ID:
+                try_mlflow_log(mlflow.start_run, _AUTOLOG_RUN_ID)
+            else:
+                try_mlflow_log(mlflow.start_run)
+                auto_end = True
+
         original = gorilla.get_original_attribute(tensorflow.estimator.Estimator,
                                                   'export_saved_model')
         serialized = original(self, *args, **kwargs)
         try_mlflow_log(log_model, tf_saved_model_dir=serialized.decode('utf-8'),
                        tf_meta_graph_tags=[tag_constants.SERVING],
                        tf_signature_def_key='predict',
                        artifact_path='model')
+        if (mlflow.active_run() is not None and mlflow.active_run().info.run_id == _AUTOLOG_RUN_ID)\
+                or auto_end:
+            try_mlflow_log(mlflow.end_run)
         return serialized
 
     @gorilla.patch(tensorflow.estimator.Estimator)
     def export_savedmodel(self, *args, **kwargs):
+        auto_end = False
+        global _AUTOLOG_RUN_ID
+        if not mlflow.active_run():
+            if _AUTOLOG_RUN_ID:
+                try_mlflow_log(mlflow.start_run, _AUTOLOG_RUN_ID)
+            else:
+                try_mlflow_log(mlflow.start_run)
+                auto_end = True
+
         original = gorilla.get_original_attribute(tensorflow.estimator.Estimator,
                                                   'export_savedmodel')
         serialized = original(self, *args, **kwargs)
         try_mlflow_log(log_model, tf_saved_model_dir=serialized.decode('utf-8'),
                        tf_meta_graph_tags=[tag_constants.SERVING],
                        tf_signature_def_key='predict',
                        artifact_path='model')
+        if (mlflow.active_run() is not None and mlflow.active_run().info.run_id == _AUTOLOG_RUN_ID)\
+                or auto_end:
+            try_mlflow_log(mlflow.end_run)
         return serialized
 
     @gorilla.patch(tensorflow.keras.Model)
     def fit(self, *args, **kwargs):
-        global _AUTO_END_RUN
-        if not mlflow.active_run():
-            try_mlflow_log(mlflow.start_run)
-            _AUTO_END_RUN = True
-
-        original = gorilla.get_original_attribute(tensorflow.keras.Model, 'fit')
-
-        # Checking if the 'callback' argument of fit() is set
-        if len(args) >= 6:
-            tmp_list = list(args)
-            tmp_list[5], log_dir = _setup_callbacks(tmp_list[5])
-            args = tuple(tmp_list)
-        elif 'callbacks' in kwargs:
-            kwargs['callbacks'], log_dir = _setup_callbacks(kwargs['callbacks'])
-        else:
-            kwargs['callbacks'], log_dir = _setup_callbacks([])
-        result = original(self, *args, **kwargs)
-        _flush_queue()
-        _log_artifacts_with_warning(local_dir=log_dir, artifact_path='tensorboard_logs')
-        shutil.rmtree(log_dir)
-
-        if _AUTO_END_RUN:
-            try_mlflow_log(mlflow.end_run)
-        _AUTO_END_RUN = False
+        with _manage_active_run():
+            original = gorilla.get_original_attribute(tensorflow.keras.Model, 'fit')
+
+            # Checking if the 'callback' argument of fit() is set
+            if len(args) >= 6:
+                tmp_list = list(args)
+                tmp_list[5], log_dir = _setup_callbacks(tmp_list[5])
+                args = tuple(tmp_list)
+            elif 'callbacks' in kwargs:
+                kwargs['callbacks'], log_dir = _setup_callbacks(kwargs['callbacks'])
+            else:
+                kwargs['callbacks'], log_dir = _setup_callbacks([])
+            result = original(self, *args, **kwargs)
+            _flush_queue()
+            _log_artifacts_with_warning(local_dir=log_dir, artifact_path='tensorboard_logs')
+            shutil.rmtree(log_dir)
 
-        return result
+            return result
 
     @gorilla.patch(tensorflow.keras.Model)
     def fit_generator(self, *args, **kwargs):
-        global _AUTO_END_RUN
-        if not mlflow.active_run():
-            try_mlflow_log(mlflow.start_run)
-            _AUTO_END_RUN = True
-
-        original = gorilla.get_original_attribute(tensorflow.keras.Model, 'fit_generator')
-
-        # Checking if the 'callback' argument of fit() is set
-        if len(args) >= 5:
-            tmp_list = list(args)
-            tmp_list[4], log_dir = _setup_callbacks(tmp_list[4])
-            args = tuple(tmp_list)
-        elif 'callbacks' in kwargs:
-            kwargs['callbacks'], log_dir = _setup_callbacks(kwargs['callbacks'])
-        else:
-            kwargs['callbacks'], log_dir = _setup_callbacks([])
-        result = original(self, *args, **kwargs)
-        _flush_queue()
-        _log_artifacts_with_warning(local_dir=log_dir, artifact_path='tensorboard_logs')
-        shutil.rmtree(log_dir)
-
-        if _AUTO_END_RUN:
-            try_mlflow_log(mlflow.end_run)
-        _AUTO_END_RUN = False
+        with _manage_active_run():
+            original = gorilla.get_original_attribute(tensorflow.keras.Model, 'fit_generator')
+
+            # Checking if the 'callback' argument of fit() is set
+            if len(args) >= 5:
+                tmp_list = list(args)
+                tmp_list[4], log_dir = _setup_callbacks(tmp_list[4])
+                args = tuple(tmp_list)
+            elif 'callbacks' in kwargs:
+                kwargs['callbacks'], log_dir = _setup_callbacks(kwargs['callbacks'])
+            else:
+                kwargs['callbacks'], log_dir = _setup_callbacks([])
+            result = original(self, *args, **kwargs)
+            _flush_queue()
+            _log_artifacts_with_warning(local_dir=log_dir, artifact_path='tensorboard_logs')
+            shutil.rmtree(log_dir)
 
-        return result
+            return result
 
     @gorilla.patch(EventFileWriter)
     def add_event(self, event):
@@ -754,6 +793,7 @@ def add_summary(self, *args, **kwargs):
     patches = [
         gorilla.Patch(EventFileWriter, 'add_event', add_event, settings=settings),
         gorilla.Patch(EventFileWriterV2, 'add_event', add_event, settings=settings),
+        gorilla.Patch(tensorflow.estimator.Estimator, 'train', train, settings=settings),
         gorilla.Patch(tensorflow.keras.Model, 'fit', fit, settings=settings),
         gorilla.Patch(tensorflow.keras.Model, 'fit_generator', fit_generator, settings=settings),
         gorilla.Patch(tensorflow.estimator.Estimator, 'export_saved_model',