Ray Tune custom callback based on model structure about lightgbm_ray HOT 2 CLOSED

class CustomTuneReportCallback(TuneReportCallback):
    def __init__(self, metrics):
        super(CustomTuneReportCallback, self).__init__(metrics)

    def __call__(self, env: CallbackEnv) -> None:
            if not self.is_rank_0:
                return
            eval_result = self._get_eval_result(env)
            report_dict = self._get_report_dict(eval_result)

            model = env.model

            mod_dump = model.dump_model()
            tree_info = mod_dump['tree_info']
            num_leaves = 0
            num_iterations = 0
            for tree in tree_info:
                num_leaves += tree['num_leaves']
                num_iterations += 1
            report_dict.update({'total_leaves': num_leaves,
                    'rmse_train': env.evaluation_result_list[0][2],
                    'num_iterations': num_iterations})

            put_queue(lambda: tune.report(**report_dict))