def record_eval_results(eval_results, metrics_logger): """ Create a callback function that records evaluation results. """ # TODO: Remove `replace("SNAPSHOT", "dev")` once the following issue is addressed: # https://github.com/dmlc/xgboost/issues/6984 from mlflow.xgboost._autolog import IS_TRAINING_CALLBACK_SUPPORTED if IS_TRAINING_CALLBACK_SUPPORTED: from mlflow.xgboost._autolog import AutologCallback # In xgboost >= 1.3.0, user-defined callbacks should inherit # `xgboost.callback.TrainingCallback`: # https://xgboost.readthedocs.io/en/latest/python/callbacks.html#defining-your-own-callback return AutologCallback(metrics_logger, eval_results) else: from mlflow.xgboost._autolog import autolog_callback return picklable_exception_safe_function( functools.partial(autolog_callback, metrics_logger=metrics_logger, eval_results=eval_results))
def test_callback_class_is_pickable(): from mlflow.xgboost._autolog import AutologCallback cb = AutologCallback(BatchMetricsLogger(run_id="1234"), eval_results={}) pickle.dumps(cb)