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
 18from typing import Optional
 19
 20from multistorageclient.types import ObjectMetadata
 21from multistorageclient.utils import calculate_worker_processes_and_threads
 22
 23from .. import StorageClient
 24from ..providers.manifest_formats import ManifestFormat
 25from ..providers.manifest_metadata import DEFAULT_MANIFEST_BASE_DIR, ManifestMetadataProvider
 26
 27

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

[docs]
 43    @staticmethod
 44    def generate_and_write_manifest(
 45        data_storage_client: StorageClient,
 46        manifest_storage_client: StorageClient,
 47        partition_keys: Optional[list[str]] = None,
 48        manifest_format: ManifestFormat = ManifestFormat.JSONL,
 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        """
 71        data_storage_provider = data_storage_client._storage_provider
 72        manifest_storage_provider = manifest_storage_client._storage_provider
 73
 74        manifest_metadata_provider = ManifestMetadataProvider(
 75            storage_provider=manifest_storage_provider, manifest_path="", writable=True, manifest_format=manifest_format
 76        )
 77
 78        if partition_keys is not None:
 79            _, num_worker_threads = calculate_worker_processes_and_threads()
 80
 81            boundaries = list(zip([""] + partition_keys, partition_keys + [None]))
 82
 83            def process_partition(boundary):
 84                start_after, end_at = boundary
 85                for object_metadata in data_storage_provider.list_objects(
 86                    path="", start_after=start_after, end_at=end_at, show_attributes=True
 87                ):
 88                    if DEFAULT_MANIFEST_BASE_DIR not in object_metadata.key.split("/"):  # Do not track manifest files
 89                        manifest_metadata_provider.add_file(path=object_metadata.key, metadata=object_metadata)
 90
 91            with ThreadPoolExecutor(max_workers=num_worker_threads) as executor:
 92                futures = [executor.submit(process_partition, boundary) for boundary in boundaries]
 93                for future in futures:
 94                    future.result()
 95        else:
 96            for object_metadata in data_storage_provider.list_objects(path="", show_attributes=True):
 97                if DEFAULT_MANIFEST_BASE_DIR not in object_metadata.key.split("/"):  # Do not track manifest files
 98                    manifest_metadata_provider.add_file(path=object_metadata.key, metadata=object_metadata)
 99
100        manifest_metadata_provider.commit_updates()