# Buggy code
def on_train_end(self, logs=None):  # pylint: disable=unused-argument
    opt = self.model.optimizer
    if hasattr(opt, "optimizer"):
        opt = opt.optimizer
    mlflow.log_param("optimizer_name", type(opt).__name__)
    if hasattr(opt, "_lr"):
        lr = (
            opt._lr
            if type(opt._lr) is float
            else tensorflow.keras.backend.eval(opt._lr)
        )
        mlflow.log_param("learning_rate", lr)
    if hasattr(opt, "_epsilon"):
        epsilon = (
            opt._epsilon
            if type(opt._epsilon) is float
            else tensorflow.keras.backend.eval(opt._epsilon)
        )
        mlflow.log_param("epsilon", epsilon)
    l = []
    self.model.summary(print_fn=l.append)
    summary = "\n".join(l)
    mlflow.set_tag("summary", summary)
    mlflow.keras.log_model(self.model, artifact_path="model")


# 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):  # pylint: disable=unused-argument
    opt = self.model.optimizer
    if hasattr(opt, "optimizer"):
        opt = opt.optimizer
        try_mlflow_log(mlflow.log_param, "optimizer_name", type(opt).__name__)
    if hasattr(opt, "_lr"):
        lr = (
            opt._lr
            if type(opt._lr) is float
            else tensorflow.keras.backend.eval(opt._lr)
        )
        try_mlflow_log(mlflow.log_param("learning_rate", lr))
    if hasattr(opt, "_epsilon"):
        epsilon = (
            opt._epsilon
            if type(opt._epsilon) is float
            else tensorflow.keras.backend.eval(opt._epsilon)
        )
        try_mlflow_log(mlflow.log_param, "epsilon", epsilon)
    l = []
    self.model.summary(print_fn=l.append)
    summary = "\n".join(l)
    try_mlflow_log(mlflow.set_tag, "summary", summary)
    try_mlflow_log(mlflow.keras.log_model, self.model, artifact_path="model")