def get_callbacks_watchlist()

in src/sagemaker_xgboost_container/algorithm_mode/train.py [0:0]

• 35 lines of code
• 12 McCabe index (conditional complexity)

def get_callbacks_watchlist(train_dmatrix, val_dmatrix, model_dir, checkpoint_dir,
                            early_stopping_data_name, early_stopping_metric, early_stopping_rounds,
                            save_model_on_termination, is_master, fold=None):
    if checkpoint_dir and fold is not None:
        checkpoint_dir = os.path.join(checkpoint_dir, f"model-{fold}")

    # Set callbacks
    xgb_model, iteration = checkpointing.load_checkpoint(checkpoint_dir)
    if xgb_model is not None:
        if fold is not None:
            xgb_model = f"{xgb_model}-{fold}"
        logging.info("Checkpoint loaded from %s", xgb_model)
        logging.info("Resuming from iteration %s", iteration)

    callbacks = []
    callbacks.append(xgb.callback.EvaluationMonitor())
    if checkpoint_dir:
        save_checkpoint = xgb.callback.TrainingCheckPoint(directory=checkpoint_dir,
                                                          iterations=iteration,
                                                          name=checkpointing.CHECKPOINT_FILENAME)
        callbacks.append(save_checkpoint)

    if save_model_on_termination == "true":
        model_name = f"{MODEL_NAME}-{fold}" if fold is not None else MODEL_NAME
        save_intermediate_model = checkpointing.SaveIntermediateModelCallBack(model_dir, model_name, is_master)
        callbacks.append(save_intermediate_model)
        add_sigterm_handler(model_dir, is_master)

    if early_stopping_data_name and early_stopping_metric and early_stopping_rounds:
        maximize = early_stopping_metric in XGB_MAXIMIZE_METRICS
        early_stop = xgb.callback.EarlyStopping(rounds=early_stopping_rounds,
                                                data_name=early_stopping_data_name,
                                                metric_name=early_stopping_metric,
                                                maximize=maximize,
                                                save_best=True)
        callbacks.append(early_stop)

    watchlist = [(train_dmatrix, 'train')]
    if val_dmatrix is not None:
        watchlist.append((val_dmatrix, 'validation'))

    return xgb_model, iteration, callbacks, watchlist