Method cache_metadata

utils/dataset.py:532–586

(self, regenerate_cache=False, trust_cache=False)

            quit()

    def cache_metadata(self, regenerate_cache=False, trust_cache=False):
        def check_grouped_metadata():
            all_grouped_metadata_exists = False
            unique_grouping_keys = None
            if self.grouping_keys_json_file.exists():
                with open(self.grouping_keys_json_file) as f:
                    unique_grouping_keys = json.load(f)
                if self.use_size_buckets and not all(len(key) == 3 for key in unique_grouping_keys):
                    # Using size buckets but have AR keys.
                    return False, unique_grouping_keys
                elif not all(len(key) == 2 for key in unique_grouping_keys):
                    # Using AR buckets but have size bucket keys
                    return False, unique_grouping_keys
                all_grouped_metadata_exists = all(
                    (self.cache_dir / f'metadata/grouped_metadata_{bucket_suffix(key)}').exists()
                    for key in unique_grouping_keys
                )
            return all_grouped_metadata_exists, unique_grouping_keys

        # Check if all the grouped metadata datasets exist. If so, we can directly load them.
        all_grouped_metadata_exists, unique_grouping_keys = check_grouped_metadata()
        if regenerate_cache or not all_grouped_metadata_exists or not trust_cache:
            # Otherwise, need to compute the ungrouped metadata and then group.
            print('Grouped metadata is not cached. Computing ungrouped metadata and then grouping.')
            unique_grouping_keys = self._group_metadata_and_save_to_disk(regenerate_cache=regenerate_cache, trust_cache=trust_cache)
        else:
            print('Found grouped metadata cache. Directly loading it.')

        for grouping_key in unique_grouping_keys:
            grouped_cache_dir = self.cache_dir / f'metadata/grouped_metadata_{bucket_suffix(grouping_key)}'
            print(f'Loading grouped metadata with grouping key {grouping_key}')
            metadata = datasets.load_from_disk(str(grouped_cache_dir))
            if self.use_size_buckets:
                assert len(grouping_key) == 3
                self.size_bucket_datasets.append(
                    SizeBucketDataset(
                        metadata,
                        self.directory_config,
                        grouping_key,
                        self.cache_dir,
                        self,
                    )
                )
            else:
                self.ar_bucket_datasets.append(
                    ARBucketDataset(
                        grouping_key,
                        self.resolutions,
                        metadata,
                        self.directory_config,
                        self.cache_dir,
                        self.round_to_multiple,
                        self,
                    )
                )

    def _group_metadata_and_save_to_disk(self, regenerate_cache=False, trust_cache=False):
        metadata_dataset = self._get_ungrouped_metadata(regenerate_cache=regenerate_cache, trust_cache=trust_cache)