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