Repository: mlflow/mlflow

Issue:
[BUG] MLflow tf.autolog() deletes TensorBoard records if both services are used.
### System information
- **Have I written custom code (as opposed to using a stock example script provided in MLflow)**: yes
- **OS Platform and Distribution (e.g., Linux Ubuntu 16.04)**: Google Colab
- **MLflow installed from (source or binary)**: source
- **MLflow version (run ``mlflow --version``)**:1.7.2
- **Python version**: 3.6.9
- **npm version, if running the dev UI**:
- **Exact command to reproduce**: https://colab.research.google.com/drive/1-bn6HGKoXEg5qGzWjvTyW34fBRgA6TJv

### Describe the problem
Whenever TensorBoard and MLflow are used in parallel, a MLflow clean up hook deletes all metrics logged to TensorBoard at the end of training. See the attached Colab for an example.

### Code to reproduce issue
https://colab.research.google.com/drive/1-bn6HGKoXEg5qGzWjvTyW34fBRgA6TJv



Discussion:

Resolving patch:
diff --git a/mlflow/tensorflow.py b/mlflow/tensorflow.py
--- a/mlflow/tensorflow.py
+++ b/mlflow/tensorflow.py
@@ -20,6 +20,7 @@
 import atexit
 import time
 import tempfile
+from collections import namedtuple
 
 import pandas
 
@@ -607,6 +608,13 @@ def _get_tensorboard_callback(lst):
     return None
 
 
+# A representation of a TensorBoard event logging directory with two attributes:
+# :location - string: The filesystem location of the logging directory
+# :is_temp - boolean: `True` if the logging directory was created for temporary use by MLflow,
+#                     `False` otherwise
+_TensorBoardLogDir = namedtuple("_TensorBoardLogDir", ["location", "is_temp"])
+
+
 def _setup_callbacks(lst):
     """
     Adds TensorBoard and MlfLowTfKeras callbacks to the
@@ -614,10 +622,10 @@ def _setup_callbacks(lst):
     """
     tb = _get_tensorboard_callback(lst)
     if tb is None:
-        log_dir = tempfile.mkdtemp()
-        out_list = lst + [TensorBoard(log_dir)]
+        log_dir = _TensorBoardLogDir(location=tempfile.mkdtemp(), is_temp=True)
+        out_list = lst + [TensorBoard(log_dir.location)]
     else:
-        log_dir = tb.log_dir
+        log_dir = _TensorBoardLogDir(location=tb.log_dir, is_temp=False)
         out_list = lst
     if LooseVersion(tensorflow.__version__) < LooseVersion('2.0.0'):
         out_list += [__MLflowTfKerasCallback()]
@@ -814,8 +822,10 @@ def fit(self, *args, **kwargs):
             _log_early_stop_callback_metrics(early_stop_callback, history)
 
             _flush_queue()
-            _log_artifacts_with_warning(local_dir=log_dir, artifact_path='tensorboard_logs')
-            shutil.rmtree(log_dir)
+            _log_artifacts_with_warning(
+                local_dir=log_dir.location, artifact_path='tensorboard_logs')
+            if log_dir.is_temp:
+                shutil.rmtree(log_dir.location)
 
             return history
 
@@ -839,8 +849,10 @@ def fit_generator(self, *args, **kwargs):
                 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)
+            _log_artifacts_with_warning(
+                local_dir=log_dir.location, artifact_path='tensorboard_logs')
+            if log_dir.is_temp:
+                shutil.rmtree(log_dir.location)
 
             return result
 


Test patch:
diff --git a/tests/tensorflow_autolog/test_tensorflow2_autolog.py b/tests/tensorflow_autolog/test_tensorflow2_autolog.py
--- a/tests/tensorflow_autolog/test_tensorflow2_autolog.py
+++ b/tests/tensorflow_autolog/test_tensorflow2_autolog.py
@@ -304,6 +304,62 @@ def test_tf_keras_autolog_non_early_stop_callback_no_log(tf_keras_random_data_ru
     assert len(metric_history) == num_of_epochs
 
 
+@pytest.mark.large
+@pytest.mark.parametrize('fit_variant', ['fit', 'fit_generator'])
+def test_tf_keras_autolog_does_not_delete_logging_directory_for_tensorboard_callback(
+        tmpdir, random_train_data, random_one_hot_labels, fit_variant):
+    tensorboard_callback_logging_dir_path = str(tmpdir.mkdir("tb_logs"))
+    tensorboard_callback = tf.keras.callbacks.TensorBoard(
+        tensorboard_callback_logging_dir_path, histogram_freq=0)
+
+    mlflow.tensorflow.autolog()
+
+    data = random_train_data
+    labels = random_one_hot_labels
+
+    model = create_tf_keras_model()
+
+    if fit_variant == 'fit_generator':
+        def generator():
+            while True:
+                yield data, labels
+        model.fit_generator(
+            generator(), epochs=10, steps_per_epoch=1, callbacks=[tensorboard_callback])
+    else:
+        model.fit(data, labels, epochs=10, callbacks=[tensorboard_callback])
+
+    assert os.path.exists(tensorboard_callback_logging_dir_path)
+
+
+@pytest.mark.large
+@pytest.mark.parametrize('fit_variant', ['fit', 'fit_generator'])
+def test_tf_keras_autolog_logs_to_and_deletes_temporary_directory_when_tensorboard_callback_absent(
+        tmpdir, random_train_data, random_one_hot_labels, fit_variant):
+    import mock
+    from mlflow.tensorflow import _TensorBoardLogDir
+
+    mlflow.tensorflow.autolog()
+
+    mock_log_dir_inst = _TensorBoardLogDir(location=str(tmpdir.mkdir("tb_logging")), is_temp=True)
+    with mock.patch("mlflow.tensorflow._TensorBoardLogDir", autospec=True) as mock_log_dir_class:
+        mock_log_dir_class.return_value = mock_log_dir_inst
+
+        data = random_train_data
+        labels = random_one_hot_labels
+
+        model = create_tf_keras_model()
+
+        if fit_variant == 'fit_generator':
+            def generator():
+                while True:
+                    yield data, labels
+            model.fit_generator(generator(), epochs=10, steps_per_epoch=1)
+        else:
+            model.fit(data, labels, epochs=10)
+
+        assert not os.path.exists(mock_log_dir_inst.location)
+
+
 def create_tf_estimator_model(dir, export):
     CSV_COLUMN_NAMES = ['SepalLength', 'SepalWidth', 'PetalLength', 'PetalWidth', 'Species']
     SPECIES = ['Setosa', 'Versicolor', 'Virginica']
@@ -399,16 +455,9 @@ def test_tf_estimator_autolog_model_can_load_from_artifact(tf_estimator_random_d
                                          "/model")
 
 
-@pytest.fixture
-def duplicate_autolog_tf_estimator_run(tmpdir, manual_run, export):
-    mlflow.tensorflow.autolog(every_n_iter=23)  # 23 is prime; no false positives in test
-    run = tf_estimator_random_data_run(tmpdir, manual_run, export)
-    return run  # should be autologged every 4 steps
-
-
 @pytest.mark.large
 @pytest.mark.parametrize('export', [True, False])
-def test_duplicate_autolog_second_overrides(duplicate_autolog_tf_estimator_run):
+def test_duplicate_autolog_second_overrides(tf_estimator_random_data_run):
     client = mlflow.tracking.MlflowClient()
-    metrics = client.get_metric_history(duplicate_autolog_tf_estimator_run.info.run_id, 'loss')
+    metrics = client.get_metric_history(tf_estimator_random_data_run.info.run_id, 'loss')
     assert all((x.step - 1) % 4 == 0 for x in metrics)
diff --git a/tests/tensorflow_autolog/test_tensorflow_autolog.py b/tests/tensorflow_autolog/test_tensorflow_autolog.py
--- a/tests/tensorflow_autolog/test_tensorflow_autolog.py
+++ b/tests/tensorflow_autolog/test_tensorflow_autolog.py
@@ -288,6 +288,62 @@ def test_tf_keras_autolog_non_early_stop_callback_no_log(tf_keras_random_data_ru
     assert len(metric_history) == num_of_epochs
 
 
+@pytest.mark.large
+@pytest.mark.parametrize('fit_variant', ['fit', 'fit_generator'])
+def test_tf_keras_autolog_does_not_delete_logging_directory_for_tensorboard_callback(
+        tmpdir, random_train_data, random_one_hot_labels, fit_variant):
+    tensorboard_callback_logging_dir_path = str(tmpdir.mkdir("tb_logs"))
+    tensorboard_callback = tf.keras.callbacks.TensorBoard(
+        tensorboard_callback_logging_dir_path, histogram_freq=0)
+
+    mlflow.tensorflow.autolog()
+
+    data = random_train_data
+    labels = random_one_hot_labels
+
+    model = create_tf_keras_model()
+
+    if fit_variant == 'fit_generator':
+        def generator():
+            while True:
+                yield data, labels
+        model.fit_generator(
+            generator(), epochs=10, steps_per_epoch=1, callbacks=[tensorboard_callback])
+    else:
+        model.fit(data, labels, epochs=10, callbacks=[tensorboard_callback])
+
+    assert os.path.exists(tensorboard_callback_logging_dir_path)
+
+
+@pytest.mark.large
+@pytest.mark.parametrize('fit_variant', ['fit', 'fit_generator'])
+def test_tf_keras_autolog_logs_to_and_deletes_temporary_directory_when_tensorboard_callback_absent(
+        tmpdir, random_train_data, random_one_hot_labels, fit_variant):
+    import mock
+    from mlflow.tensorflow import _TensorBoardLogDir
+
+    mlflow.tensorflow.autolog()
+
+    mock_log_dir_inst = _TensorBoardLogDir(location=str(tmpdir.mkdir("tb_logging")), is_temp=True)
+    with mock.patch("mlflow.tensorflow._TensorBoardLogDir", autospec=True) as mock_log_dir_class:
+        mock_log_dir_class.return_value = mock_log_dir_inst
+
+        data = random_train_data
+        labels = random_one_hot_labels
+
+        model = create_tf_keras_model()
+
+        if fit_variant == 'fit_generator':
+            def generator():
+                while True:
+                    yield data, labels
+            model.fit_generator(generator(), epochs=10, steps_per_epoch=1)
+        else:
+            model.fit(data, labels, epochs=10)
+
+        assert not os.path.exists(mock_log_dir_inst.location)
+
+
 @pytest.fixture
 def tf_core_random_tensors():
     mlflow.tensorflow.autolog(every_n_iter=4)
@@ -423,16 +479,10 @@ def test_tf_estimator_autolog_model_can_load_from_artifact(tf_estimator_random_d
                                          "/model", session)
 
 
-@pytest.fixture
-def duplicate_autolog_tf_estimator_run(tmpdir, manual_run, export):
-    mlflow.tensorflow.autolog(every_n_iter=23)  # 23 is prime; no false positives in test
-    run = tf_estimator_random_data_run(tmpdir, manual_run, export)
-    return run  # should be autologged every 4 steps
-
-
 @pytest.mark.large
 @pytest.mark.parametrize('export', [True, False])
-def test_duplicate_autolog_second_overrides(duplicate_autolog_tf_estimator_run):
+def test_duplicate_autolog_second_overrides(tf_estimator_random_data_run):
+    mlflow.tensorflow.autolog(every_n_iter=23)
     client = mlflow.tracking.MlflowClient()
-    metrics = client.get_metric_history(duplicate_autolog_tf_estimator_run.info.run_id, 'loss')
+    metrics = client.get_metric_history(tf_estimator_random_data_run.info.run_id, 'loss')
     assert all((x.step - 1) % 4 == 0 for x in metrics)