# Buggy code
def autolog(every_n_iter=100):
    # pylint: disable=E0611
    """
    Enable automatic logging from TensorFlow to MLflow. If applicable,
    model checkpoints are logged as artifacts to a 'models' directory, along
    with any TensorBoard log data.

    Refer to the tracking documentation for
    information on what is logged with different TensorFlow workflows.

    :param every_n_iter: The frequency with which metrics should be logged.
                                  Defaults to 100. Ex: a value of 100 will log metrics
                                  at step 0, 100, 200, etc.

    """
    global _LOG_EVERY_N_STEPS
    _LOG_EVERY_N_STEPS = every_n_iter

    from distutils.version import StrictVersion

    if StrictVersion(tensorflow.__version__) < StrictVersion("1.12") or StrictVersion(
        tensorflow.__version__
    ) >= StrictVersion("2.0"):
        warnings.warn(
            "Could not log to MLflow. Only TensorFlow versions"
            + "1.12 <= v < 2.0.0 are supported."
        )
        return

    try:
        from tensorflow.python.summary.writer.event_file_writer import EventFileWriter
        from tensorflow.python.summary.writer.event_file_writer_v2 import (
            EventFileWriterV2,
        )
        from tensorflow.python.saved_model import tag_constants
        from tensorflow.python.summary.writer.writer import FileWriter
    except ImportError:
        warnings.warn(
            "Could not log to MLflow. Only TensorFlow versions"
            + "1.12 <= v < 2.0.0 are supported."
        )
        return

    @gorilla.patch(tensorflow.estimator.Estimator)
    def export_saved_model(self, *args, **kwargs):
        original = gorilla.get_original_attribute(
            tensorflow.estimator.Estimator, "export_saved_model"
        )
        serialized = original(self, *args, **kwargs)
        try:
            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",
            )
        except MlflowException as e:
            warnings.warn("Logging to MLflow failed: " + str(e))
        return serialized

    @gorilla.patch(tensorflow.estimator.Estimator)
    def export_savedmodel(self, *args, **kwargs):
        original = gorilla.get_original_attribute(
            tensorflow.estimator.Estimator, "export_savedmodel"
        )
        serialized = original(self, *args, **kwargs)
        try:
            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",
            )
        except MlflowException as e:
            warnings.warn("Logging to MLflow failed: " + str(e))
        return serialized

    @gorilla.patch(tensorflow.keras.Model)
    def fit(self, *args, **kwargs):
        original = gorilla.get_original_attribute(tensorflow.keras.Model, "fit")
        if len(args) >= 6:
            l = list(args)
            l[5], log_dir = _setup_callbacks(l[5])
            args = tuple(l)
        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

    @gorilla.patch(EventFileWriter)
    def add_event(self, event):
        _log_event(event)
        original = gorilla.get_original_attribute(EventFileWriter, "add_event")
        return original(self, event)

    @gorilla.patch(FileWriter)
    def add_summary(self, *args, **kwargs):
        original = gorilla.get_original_attribute(FileWriter, "add_summary")
        result = original(self, *args, **kwargs)
        _flush_queue()
        return result

    settings = gorilla.Settings(allow_hit=True, store_hit=True)
    patches = [
        gorilla.Patch(EventFileWriter, "add_event", add_event, settings=settings),
        gorilla.Patch(EventFileWriterV2, "add_event", add_event, settings=settings),
        gorilla.Patch(tensorflow.keras.Model, "fit", fit, settings=settings),
        gorilla.Patch(
            tensorflow.estimator.Estimator,
            "export_saved_model",
            export_saved_model,
            settings=settings,
        ),
        gorilla.Patch(
            tensorflow.estimator.Estimator,
            "export_savedmodel",
            export_savedmodel,
            settings=settings,
        ),
        gorilla.Patch(FileWriter, "add_summary", add_summary, settings=settings),
    ]

    for x in patches:
        gorilla.apply(x)


# Traceback (sqlalchemy.exc.DataError)
2019/08/01 18:01:53 ERROR mlflow.server: Exception on /api/2.0/mlflow/runs/set-tag [POST]
Traceback (most recent call last):
  File "/home/myserver/mlflow_env/lib/python3.6/site-packages/sqlalchemy/engine/base.py", line 1244, in _execute_context
    cursor, statement, parameters, context
  File "/home/myserver/mlflow_env/lib/python3.6/site-packages/sqlalchemy/engine/default.py", line 552, in do_execute
    cursor.execute(statement, parameters)
psycopg2.errors.StringDataRightTruncation: value too long for type character varying(250)


The above exception was the direct cause of the following exception:

Traceback (most recent call last):
  File "/home/myserver/mlflow_env/lib/python3.6/site-packages/mlflow/store/sqlalchemy_store.py", line 148, in make_managed_session
    session.commit()
  File "/home/myserver/mlflow_env/lib/python3.6/site-packages/sqlalchemy/orm/session.py", line 1027, in commit
    self.transaction.commit()
  File "/home/myserver/mlflow_env/lib/python3.6/site-packages/sqlalchemy/orm/session.py", line 494, in commit
    self._prepare_impl()
  File "/home/myserver/mlflow_env/lib/python3.6/site-packages/sqlalchemy/orm/session.py", line 473, in _prepare_impl
    self.session.flush()
  File "/home/myserver/mlflow_env/lib/python3.6/site-packages/sqlalchemy/orm/session.py", line 2459, in flush
    self._flush(objects)
  File "/home/myserver/mlflow_env/lib/python3.6/site-packages/sqlalchemy/orm/session.py", line 2597, in _flush
    transaction.rollback(_capture_exception=True)
  File "/home/myserver/mlflow_env/lib/python3.6/site-packages/sqlalchemy/util/langhelpers.py", line 68, in __exit__
    compat.reraise(exc_type, exc_value, exc_tb)
  File "/home/myserver/mlflow_env/lib/python3.6/site-packages/sqlalchemy/util/compat.py", line 153, in reraise
    raise value
  File "/home/myserver/mlflow_env/lib/python3.6/site-packages/sqlalchemy/orm/session.py", line 2557, in _flush
    flush_context.execute()
  File "/home/myserver/mlflow_env/lib/python3.6/site-packages/sqlalchemy/orm/unitofwork.py", line 422, in execute
    rec.execute(self)
  File "/home/myserver/mlflow_env/lib/python3.6/site-packages/sqlalchemy/orm/unitofwork.py", line 589, in execute
    uow,
  File "/home/myserver/mlflow_env/lib/python3.6/site-packages/sqlalchemy/orm/persistence.py", line 245, in save_obj
    insert,
  File "/home/myserver/mlflow_env/lib/python3.6/site-packages/sqlalchemy/orm/persistence.py", line 1084, in _emit_insert_statements
    c = cached_connections[connection].execute(statement, multiparams)
  File "/home/myserver/mlflow_env/lib/python3.6/site-packages/sqlalchemy/engine/base.py", line 988, in execute
    return meth(self, multiparams, params)
  File "/home/myserver/mlflow_env/lib/python3.6/site-packages/sqlalchemy/sql/elements.py", line 287, in _execute_on_connection
    return connection._execute_clauseelement(self, multiparams, params)
  File "/home/myserver/mlflow_env/lib/python3.6/site-packages/sqlalchemy/engine/base.py", line 1107, in _execute_clauseelement
    distilled_params,
  File "/home/myserver/mlflow_env/lib/python3.6/site-packages/sqlalchemy/engine/base.py", line 1248, in _execute_context
    e, statement, parameters, cursor, context
  File "/home/myserver/mlflow_env/lib/python3.6/site-packages/sqlalchemy/engine/base.py", line 1466, in _handle_dbapi_exception
    util.raise_from_cause(sqlalchemy_exception, exc_info)
  File "/home/myserver/mlflow_env/lib/python3.6/site-packages/sqlalchemy/util/compat.py", line 398, in raise_from_cause
    reraise(type(exception), exception, tb=exc_tb, cause=cause)
  File "/home/myserver/mlflow_env/lib/python3.6/site-packages/sqlalchemy/util/compat.py", line 152, in reraise
    raise value.with_traceback(tb)
  File "/home/myserver/mlflow_env/lib/python3.6/site-packages/sqlalchemy/engine/base.py", line 1244, in _execute_context
    cursor, statement, parameters, context
  File "/home/myserver/mlflow_env/lib/python3.6/site-packages/sqlalchemy/engine/default.py", line 552, in do_execute
    cursor.execute(statement, parameters)
sqlalchemy.exc.DataError: (psycopg2.errors.StringDataRightTruncation) value too long for type character varying(250)

[SQL: INSERT INTO tags (key, value, run_uuid) VALUES (%(key)s, %(value)s, %(run_uuid)s)]
[parameters: {'key': 'summary', 'value': '_________________________________________________________________\nLayer (type)                 Output Shape              Param #   \n=============== ... (713 characters truncated) ... =======\nTotal params: 536,110\nTrainable params: 536,110\nNon-trainable params: 0\n_________________________________________________________________', 'run_uuid': ‘123f3e6ce63d41e9ba51644694b4a37f'}]
(Background on this error at: http://sqlalche.me/e/9h9h)


# Fixed code
def autolog(every_n_iter=100):
    # pylint: disable=E0611
    """
    Enable automatic logging from TensorFlow to MLflow. If applicable,
    model checkpoints are logged as artifacts to a 'models' directory, along
    with any TensorBoard log data.

    Refer to the tracking documentation for
    information on what is logged with different TensorFlow workflows.

    :param every_n_iter: The frequency with which metrics should be logged.
                                  Defaults to 100. Ex: a value of 100 will log metrics
                                  at step 0, 100, 200, etc.

    """
    global _LOG_EVERY_N_STEPS
    _LOG_EVERY_N_STEPS = every_n_iter

    from distutils.version import StrictVersion

    if StrictVersion(tensorflow.__version__) < StrictVersion("1.12") or StrictVersion(
        tensorflow.__version__
    ) >= StrictVersion("2.0"):
        warnings.warn(
            "Could not log to MLflow. Only TensorFlow versions"
            + "1.12 <= v < 2.0.0 are supported."
        )
        return

    try:
        from tensorflow.python.summary.writer.event_file_writer import EventFileWriter
        from tensorflow.python.summary.writer.event_file_writer_v2 import (
            EventFileWriterV2,
        )
        from tensorflow.python.saved_model import tag_constants
        from tensorflow.python.summary.writer.writer import FileWriter
    except ImportError:
        warnings.warn(
            "Could not log to MLflow. Only TensorFlow versions"
            + "1.12 <= v < 2.0.0 are supported."
        )
        return

    @gorilla.patch(tensorflow.estimator.Estimator)
    def export_saved_model(self, *args, **kwargs):
        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",
        )
        return serialized

    @gorilla.patch(tensorflow.estimator.Estimator)
    def export_savedmodel(self, *args, **kwargs):
        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",
        )
        return serialized

    @gorilla.patch(tensorflow.keras.Model)
    def fit(self, *args, **kwargs):
        original = gorilla.get_original_attribute(tensorflow.keras.Model, "fit")
        if len(args) >= 6:
            l = list(args)
            l[5], log_dir = _setup_callbacks(l[5])
            args = tuple(l)
        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

    @gorilla.patch(EventFileWriter)
    def add_event(self, event):
        _log_event(event)
        original = gorilla.get_original_attribute(EventFileWriter, "add_event")
        return original(self, event)

    @gorilla.patch(FileWriter)
    def add_summary(self, *args, **kwargs):
        original = gorilla.get_original_attribute(FileWriter, "add_summary")
        result = original(self, *args, **kwargs)
        _flush_queue()
        return result

    settings = gorilla.Settings(allow_hit=True, store_hit=True)
    patches = [
        gorilla.Patch(EventFileWriter, "add_event", add_event, settings=settings),
        gorilla.Patch(EventFileWriterV2, "add_event", add_event, settings=settings),
        gorilla.Patch(tensorflow.keras.Model, "fit", fit, settings=settings),
        gorilla.Patch(
            tensorflow.estimator.Estimator,
            "export_saved_model",
            export_saved_model,
            settings=settings,
        ),
        gorilla.Patch(
            tensorflow.estimator.Estimator,
            "export_savedmodel",
            export_savedmodel,
            settings=settings,
        ),
        gorilla.Patch(FileWriter, "add_summary", add_summary, settings=settings),
    ]

    for x in patches:
        gorilla.apply(x)