def test_callback_is_picklable():
cb = __MLflowTfKeras2Callback(
log_models=True,
metrics_logger=BatchMetricsLogger(run_id="1234"),
log_every_n_steps=5)
pickle.dumps(cb)
tb = _TensorBoard()
pickle.dumps(tb)
def _setup_callbacks(lst, log_models, metrics_logger):
"""
Adds TensorBoard and MlfLowTfKeras callbacks to the
input list, and returns the new list and appropriate log directory.
"""
# pylint: disable=no-name-in-module
from mlflow.tensorflow._autolog import _TensorBoard, __MLflowTfKeras2Callback
tb = _get_tensorboard_callback(lst)
if tb is None:
log_dir = _TensorBoardLogDir(location=tempfile.mkdtemp(), is_temp=True)
out_list = lst + [_TensorBoard(log_dir.location)]
else:
log_dir = _TensorBoardLogDir(location=tb.log_dir, is_temp=False)
out_list = lst
out_list += [__MLflowTfKeras2Callback(log_models, metrics_logger, _LOG_EVERY_N_STEPS)]
return out_list, log_dir