# Buggy code
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))


# 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 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")