def _prepare_collections()

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

• 28 lines of code
• 16 McCabe index (conditional complexity)

    def _prepare_collections(self):
        """Populate collections_to_save and ensure every collection has
        a save_config and reduction_config."""
        for c_name, c in self.collection_manager.get_collections().items():
            if c_name not in self._get_default_collections():
                if bool(c.include_regex) is False and bool(c.tensor_names) is False:
                    raise InvalidCollectionConfiguration(c_name)
            if c in self._collections_to_save:
                continue
            elif self._should_collection_be_saved(CollectionKeys.ALL):
                self._collections_to_save.add(c)
            elif self._should_collection_be_saved(c_name):
                self._collections_to_save.add(c)

        self.logger.info(
            f'Monitoring the collections: {", ".join([x.name for x in self._collections_to_save])}'
        )
        # Populate configs_for_collections and reduction_config
        for c_name, c in self.collection_manager.get_collections().items():

            if c_name in NON_HISTOGRAM_COLLECTIONS:
                c.save_histogram = False

            if c.save_config is None:
                # Set to the default if None
                c.save_config = self.save_config
            elif isinstance(c.save_config, SaveConfig):
                # Otherwise, set missing modes to the defaults
                c.save_config.merge_default_save_config(self.save_config)
            else:
                raise TypeError(f"save_config={c.save_config} must be None or SaveConfig")

            if c_name in NON_REDUCTION_COLLECTIONS:
                c.reduction_config = ReductionConfig(save_raw_tensor=True)
            elif c.reduction_config is None:
                c.reduction_config = self.reduction_config

        self.prepared_collections = True