# Buggy code
def autolog():
    """
    Enable automatic logging from TensorFlow to MLflow.
    Logs loss and any other metrics specified in the fit
    function, and optimizer data as parameters. Model checkpoints
    are logged as artifacts to a 'models' directory.
    """
    import keras

    class __MLflowKerasCallback(keras.callbacks.Callback):
        """
        Callback for auto-logging metrics and parameters.
        Records available logs after each epoch.
        Records model structural information as params after training finishes.
        """

        def on_epoch_end(self, epoch, logs=None):
            if not logs:
                return
            try:
                mlflow.log_metrics(logs, step=epoch)
            except mlflow.exceptions.MlflowException as e:
                warnings.warn("Logging to MLflow failed: " + str(e))

        def on_train_end(self, logs=None):
            try:
                mlflow.log_param("num_layers", len(self.model.layers))
                mlflow.log_param("optimizer_name", type(self.model.optimizer).__name__)
                if hasattr(self.model.optimizer, "lr"):
                    lr = (
                        self.model.optimizer.lr
                        if type(self.model.optimizer.lr) is float
                        else keras.backend.eval(self.model.optimizer.lr)
                    )
                    mlflow.log_param("learning_rate", lr)
                if hasattr(self.model.optimizer, "epsilon"):
                    epsilon = (
                        self.model.optimizer.epsilon
                        if type(self.model.optimizer.epsilon) is float
                        else keras.backend.eval(self.model.optimizer.epsilon)
                    )
                    mlflow.log_param("epsilon", epsilon)
                sum_list = []
                self.model.summary(print_fn=sum_list.append)
                summary = "\n".join(sum_list)
                mlflow.set_tag("summary", summary)
                log_model(self.model, artifact_path="model")
            except mlflow.exceptions.MlflowException as e:
                warnings.warn("Logging to Mlflow failed: " + str(e))

    @gorilla.patch(keras.Model)
    def fit(self, *args, **kwargs):
        original = gorilla.get_original_attribute(keras.Model, "fit")
        if len(args) >= 6:
            l = list(args)
            l[5] += [__MLflowKerasCallback()]
            args = tuple(l)
        elif "callbacks" in kwargs:
            kwargs["callbacks"] += [__MLflowKerasCallback()]
        else:
            kwargs["callbacks"] = [__MLflowKerasCallback()]
        return original(self, *args, **kwargs)

    settings = gorilla.Settings(allow_hit=True, store_hit=True)
    patch = gorilla.Patch(keras.Model, "fit", fit, settings=settings)
    gorilla.apply(patch)


# 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():
    """
    Enable automatic logging from TensorFlow to MLflow.
    Logs loss and any other metrics specified in the fit
    function, and optimizer data as parameters. Model checkpoints
    are logged as artifacts to a 'models' directory.
    """
    import keras

    class __MLflowKerasCallback(keras.callbacks.Callback):
        """
        Callback for auto-logging metrics and parameters.
        Records available logs after each epoch.
        Records model structural information as params after training finishes.
        """

        def on_epoch_end(self, epoch, logs=None):
            if not logs:
                return
            try_mlflow_log(mlflow.log_metrics, logs, step=epoch)

        def on_train_end(self, logs=None):
            try_mlflow_log(mlflow.log_param, "num_layers", len(self.model.layers))
            try_mlflow_log(
                mlflow.log_param, "optimizer_name", type(self.model.optimizer).__name__
            )
            if hasattr(self.model.optimizer, "lr"):
                lr = (
                    self.model.optimizer.lr
                    if type(self.model.optimizer.lr) is float
                    else keras.backend.eval(self.model.optimizer.lr)
                )
                try_mlflow_log(mlflow.log_param, "learning_rate", lr)
            if hasattr(self.model.optimizer, "epsilon"):
                epsilon = (
                    self.model.optimizer.epsilon
                    if type(self.model.optimizer.epsilon) is float
                    else keras.backend.eval(self.model.optimizer.epsilon)
                )
                try_mlflow_log(mlflow.log_param, "epsilon", epsilon)
            sum_list = []
            self.model.summary(print_fn=sum_list.append)
            summary = "\n".join(sum_list)
            try_mlflow_log(mlflow.set_tag, "summary", summary)
            try_mlflow_log(log_model, self.model, artifact_path="model")

    @gorilla.patch(keras.Model)
    def fit(self, *args, **kwargs):
        original = gorilla.get_original_attribute(keras.Model, "fit")
        if len(args) >= 6:
            l = list(args)
            l[5] += [__MLflowKerasCallback()]
            args = tuple(l)
        elif "callbacks" in kwargs:
            kwargs["callbacks"] += [__MLflowKerasCallback()]
        else:
            kwargs["callbacks"] = [__MLflowKerasCallback()]
        return original(self, *args, **kwargs)

    settings = gorilla.Settings(allow_hit=True, store_hit=True)
    patch = gorilla.Patch(keras.Model, "fit", fit, settings=settings)
    gorilla.apply(patch)