def _callback()

in smdebug/xgboost/hook.py [0:0]

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

    def _callback(self, env: CallbackEnv) -> None:
        if not self.prepared_collections:
            # at this point we need all collections to be ready
            # this may not be the case at creation of hook
            # as user's code after hook might add collections
            self._prepare_collections()
            self.prepared_collections = True

        self._increment_step(env.iteration)

        if self.last_saved_step is not None and not self.exported_collections:
            self.export_collections()
            self.exported_collections = True

        if not self._get_collections_to_save_for_step():
            self.logger.debug("Skipping iteration {}".format(self.step))
            return

        self._initialize_writers()

        if self._is_collection_being_saved_for_step(CollectionKeys.HYPERPARAMETERS):
            self.write_hyperparameters(env)

        if self._is_collection_being_saved_for_step(CollectionKeys.METRICS):
            self.write_metrics(env)

        if self._is_collection_being_saved_for_step(CollectionKeys.LOSSES):
            self.write_losses(env)

        if self._is_collection_being_saved_for_step(CollectionKeys.PREDICTIONS):
            self.write_predictions(env)

        if self._is_collection_being_saved_for_step(CollectionKeys.LABELS):
            self.write_labels(env)

        if self._is_collection_being_saved_for_step(CollectionKeys.FEATURE_IMPORTANCE):
            self.write_feature_importances(env)

        if self._is_collection_being_saved_for_step(CollectionKeys.TREES):
            self.write_tree_model(env)

        if self._is_collection_being_saved_for_step(CollectionKeys.FULL_SHAP):
            self._maybe_compute_shap_values(env)
            self.write_full_shap(env)

        if self._is_collection_being_saved_for_step(CollectionKeys.AVERAGE_SHAP):
            self._maybe_compute_shap_values(env)
            self.write_average_shap(env)

        self._clear_shap_values()
        self.last_saved_step = self.step

        self._close_writers()