Source code for multistorageclient.generators.manifest_metadata

  1# SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
  2# SPDX-License-Identifier: Apache-2.0
  3#
  4# Licensed under the Apache License, Version 2.0 (the "License");
  5# you may not use this file except in compliance with the License.
  6# You may obtain a copy of the License at
  7#
  8# http://www.apache.org/licenses/LICENSE-2.0
  9#
 10# Unless required by applicable law or agreed to in writing, software
 11# distributed under the License is distributed on an "AS IS" BASIS,
 12# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 13# See the License for the specific language governing permissions and
 14# limitations under the License.
 15
 16import json
 17from concurrent.futures import ThreadPoolExecutor
 18
 19from multistorageclient.types import ObjectMetadata
 20from multistorageclient.utils import calculate_worker_processes_and_threads
 21
 22from .. import StorageClient
 23from ..providers.manifest_formats import ManifestFormat
 24from ..providers.manifest_metadata import DEFAULT_MANIFEST_BASE_DIR, ManifestMetadataProvider
 25
 26

[docs]
 27class ManifestMetadataGenerator:
 28    """
 29    Generates a file metadata manifest for use with a :py:class:`multistorageclient.providers.ManifestMetadataProvider`.
 30    """
 31
 32    @staticmethod
 33    def _generate_manifest_part_body(object_metadata: list[ObjectMetadata]) -> bytes:
 34        return "\n".join(
 35            [
 36                json.dumps({**metadata_dict, "size_bytes": metadata_dict.pop("content_length")})
 37                for metadata in object_metadata
 38                for metadata_dict in [metadata.to_dict()]
 39            ]
 40        ).encode(encoding="utf-8")
 41

[docs]
 42    @staticmethod
 43    def generate_and_write_manifest(
 44        data_storage_client: StorageClient,
 45        manifest_storage_client: StorageClient,
 46        partition_keys: list[str] | None = None,
 47        manifest_format: ManifestFormat = ManifestFormat.JSONL,
 48        allow_overwrites: bool = True,
 49    ) -> None:
 50        """
 51        Generates a file metadata manifest.
 52
 53        The data storage client's base path should be set to the root path for data objects (e.g. ``my-bucket/my-data-prefix``).
 54
 55        The manifest storage client's base path should be set to the root path for manifest objects (e.g. ``my-bucket/my-manifest-prefix``).
 56
 57        The following manifest objects will be written with the destination storage client (with the total number of manifest parts being variable)::
 58
 59           .msc_manifests/
 60           ├── msc_manifest_index.json
 61           └── parts/
 62               ├── msc_manifest_part000001.jsonl (or .parquet)
 63               ├── ...
 64               └── msc_manifest_part999999.jsonl (or .parquet)
 65
 66        :param data_storage_client: Storage client for reading data objects.
 67        :param manifest_storage_client: Storage client for writing manifest objects.
 68        :param partition_keys: Optional list of keys to partition the listing operation. If provided, objects will be listed concurrently using these keys as boundaries.
 69        :param manifest_format: Format for manifest parts. Defaults to ManifestFormat.JSONL.
 70        :param allow_overwrites: Whether to allow overwriting existing files in the manifest. Defaults to True for backwards compatibility.
 71        """
 72        data_storage_provider = data_storage_client._storage_provider
 73        manifest_storage_provider = manifest_storage_client._storage_provider
 74
 75        if data_storage_provider is None or manifest_storage_provider is None:
 76            raise ValueError(
 77                "ManifestMetadataGenerator requires SingleStorageClient with _storage_provider. "
 78                "CompositeStorageClient (multi-backend) is not supported."
 79            )
 80
 81        manifest_metadata_provider = ManifestMetadataProvider(
 82            storage_provider=manifest_storage_provider,
 83            manifest_path="",
 84            writable=True,
 85            manifest_format=manifest_format,
 86            allow_overwrites=allow_overwrites,
 87        )
 88
 89        if partition_keys is not None:
 90            _, num_worker_threads = calculate_worker_processes_and_threads()
 91
 92            boundaries = list(zip([""] + partition_keys, partition_keys + [None]))
 93
 94            def process_partition(boundary):
 95                start_after, end_at = boundary
 96                for object_metadata in data_storage_provider.list_objects(
 97                    path="", start_after=start_after, end_at=end_at, show_attributes=True
 98                ):
 99                    if DEFAULT_MANIFEST_BASE_DIR not in object_metadata.key.split("/"):  # Do not track manifest files
100                        manifest_metadata_provider.add_file(path=object_metadata.key, metadata=object_metadata)
101
102            with ThreadPoolExecutor(max_workers=num_worker_threads) as executor:
103                futures = [executor.submit(process_partition, boundary) for boundary in boundaries]
104                for future in futures:
105                    future.result()
106        else:
107            for object_metadata in data_storage_provider.list_objects(path="", show_attributes=True):
108                if DEFAULT_MANIFEST_BASE_DIR not in object_metadata.key.split("/"):  # Do not track manifest files
109                    manifest_metadata_provider.add_file(path=object_metadata.key, metadata=object_metadata)
110
111        manifest_metadata_provider.commit_updates()