Source code for multistorageclient.providers.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
16from __future__ import annotations # Enables forward references in type hints
17
18import io
19import json
20import logging
21import os
22from collections.abc import Iterator
23from dataclasses import asdict, dataclass
24from datetime import datetime, timezone
25from typing import Any, Optional
26
27from ..types import MetadataProvider, ObjectMetadata, StorageProvider
28from ..utils import glob
29
30logger = logging.Logger(__name__)
31
32
33DEFAULT_MANIFEST_BASE_DIR = ".msc_manifests"
34MANIFEST_INDEX_FILENAME = "msc_manifest_index.json"
35MANIFEST_PARTS_CHILD_DIR = "parts"
36MANIFEST_PART_PREFIX = "msc_manifest_part"
37MANIFEST_PART_SUFFIX = ".jsonl" # Suffix for the manifest part files
38SEQUENCE_PADDING = 6 # Define padding for the sequence number (e.g., 6 for "000001")
39
40
[docs]
41@dataclass
42class ManifestPartReference:
43 """
44 A data class representing a reference to dataset manifest part.
45 """
46
47 #: The path of the manifest part relative to the main manifest.
48 path: str
49
[docs]
50 @staticmethod
51 def from_dict(data: dict[str, Any]) -> ManifestPartReference:
52 """
53 Creates a ManifestPartReference instance from a dictionary.
54 """
55 # Validate that the required 'path' field is present
56 if "path" not in data:
57 raise ValueError("Missing required field: 'path'")
58
59 return ManifestPartReference(path=data["path"])
60
[docs]
61 def to_dict(self) -> dict:
62 """
63 Converts ManifestPartReference instance to a dictionary.
64 """
65 return {
66 "path": self.path,
67 }
68
69
[docs]
70@dataclass
71class Manifest:
72 """
73 A data class representing a dataset manifest.
74 """
75
76 #: Defines the version of the manifest schema.
77 version: str
78 #: References to manifest parts.
79 parts: list[ManifestPartReference]
80
[docs]
81 @staticmethod
82 def from_dict(data: dict) -> "Manifest":
83 """
84 Creates a Manifest instance from a dictionary (parsed from JSON).
85 """
86 # Perform any necessary validation here
87 try:
88 version = data["version"]
89 parts = [ManifestPartReference.from_dict(part) for part in data["parts"]]
90 except KeyError as e:
91 raise ValueError("Invalid manifest data: Missing required field") from e
92
93 return Manifest(version=version, parts=parts)
94
[docs]
95 def to_json(self) -> str:
96 # Convert dataclass to dict and parts to JSON-compatible format
97 data = asdict(self)
98 data["parts"] = [part.to_dict() for part in self.parts]
99 return json.dumps(data)
100
101
102def _metadata_to_manifest_dict(metadata: ObjectMetadata) -> dict:
103 """
104 Convert an ObjectMetadata instance to a dictionary suitable with manifest format,
105 replacing 'content_length' with 'size_bytes' and removing 'content_length'.
106 """
107 metadata_dict = metadata.to_dict()
108 # Pop out content_length, store it in size_bytes
109 size_bytes = metadata_dict.pop("content_length", None)
110 metadata_dict["size_bytes"] = size_bytes
111 return metadata_dict
112
113
[docs]
114class ManifestMetadataProvider(MetadataProvider):
115 _storage_provider: StorageProvider
116 _files: dict[str, ObjectMetadata]
117 _pending_adds: dict[str, ObjectMetadata]
118 _pending_removes: set[str]
119 _manifest_path: str
120 _writable: bool
121
122 def __init__(self, storage_provider: StorageProvider, manifest_path: str, writable: bool = False) -> None:
123 """
124 Creates a :py:class:`ManifestMetadataProvider`.
125
126 :param storage_provider: Storage provider.
127 :param manifest_path: Main manifest file path.
128 :param writable: If true, allows modifications and new manifests to be written.
129 """
130 self._storage_provider = storage_provider
131 self._files = {}
132 self._pending_adds = {}
133 self._pending_removes = set()
134 self._manifest_path = manifest_path
135 self._writable = writable
136
137 self._load_manifest(storage_provider, self._manifest_path)
138
139 def _load_manifest(self, storage_provider: StorageProvider, manifest_path: str) -> None:
140 """
141 Loads manifest.
142
143 :param storage_provider: Storage provider.
144 :param manifest_path: Main manifest file path
145 """
146
147 def helper_find_manifest_file(manifest_path: str) -> str:
148 if storage_provider.is_file(manifest_path):
149 return manifest_path
150
151 if storage_provider.is_file(os.path.join(manifest_path, MANIFEST_INDEX_FILENAME)):
152 return os.path.join(manifest_path, MANIFEST_INDEX_FILENAME)
153
154 # Now go looking and select newest manifest.
155 if DEFAULT_MANIFEST_BASE_DIR not in manifest_path.split("/"):
156 manifest_path = os.path.join(manifest_path, DEFAULT_MANIFEST_BASE_DIR)
157
158 candidates = storage_provider.glob(os.path.join(manifest_path, "*", MANIFEST_INDEX_FILENAME))
159 candidates = sorted(candidates)
160 return candidates[-1] if candidates else ""
161
162 resolved_manifest_path = helper_find_manifest_file(manifest_path)
163 if not resolved_manifest_path:
164 logger.warning(f"No manifest found at '{manifest_path}'.")
165 return
166
167 file_content = storage_provider.get_object(resolved_manifest_path)
168
169 prefix = os.path.dirname(resolved_manifest_path)
170 _, file_extension = os.path.splitext(resolved_manifest_path)
171 self._load_manifest_file(storage_provider, file_content, prefix, file_extension[1:])
172
173 def _load_manifest_file(
174 self, storage_provider: StorageProvider, file_content: bytes, manifest_base: str, file_type: str
175 ) -> None:
176 """
177 Loads a manifest.
178
179 :param storage_provider: Storage provider.
180 :param file_content: Manifest file content bytes.
181 :param manifest_base: Manifest file base path.
182 :param file_type: Manifest file type.
183 """
184 if file_type == "json":
185 manifest_dict = json.loads(file_content.decode("utf-8"))
186 manifest = Manifest.from_dict(manifest_dict)
187
188 # Check manifest version. Not needed once we make the manifest model use sum types/discriminated unions.
189 if manifest.version != "1":
190 raise ValueError(f"Manifest version {manifest.version} is not supported.")
191
192 # Load manifest parts.
193 for manifest_part_reference in manifest.parts:
194 object_metadata: list[ObjectMetadata] = self._load_manifest_part_file(
195 storage_provider=storage_provider,
196 manifest_base=manifest_base,
197 manifest_part_reference=manifest_part_reference,
198 )
199
200 for object_metadatum in object_metadata:
201 self._files[object_metadatum.key] = object_metadatum
202 else:
203 raise NotImplementedError(f"Manifest file type {file_type} is not supported.")
204
205 def _load_manifest_part_file(
206 self, storage_provider: StorageProvider, manifest_base: str, manifest_part_reference: ManifestPartReference
207 ) -> list[ObjectMetadata]:
208 """
209 Loads a manifest part.
210
211 :param storage_provider: Storage provider.
212 :param manifest_base: Manifest file base path. Prepend to manifest part reference paths.
213 :param manifest_part_reference: Manifest part reference.
214 """
215 object_metadata = []
216
217 if not os.path.isabs(manifest_part_reference.path):
218 remote_path = os.path.join(manifest_base, manifest_part_reference.path)
219 else:
220 remote_path = manifest_part_reference.path
221 manifest_part_file_content = storage_provider.get_object(remote_path)
222
223 # The manifest part is a JSON lines file. Each line is a JSON-serialized ObjectMetadata.
224 for line in io.TextIOWrapper(io.BytesIO(manifest_part_file_content), encoding="utf-8"):
225 object_metadatum_dict = json.loads(line)
226 object_metadatum_dict["content_length"] = object_metadatum_dict.pop("size_bytes")
227 object_metadatum = ObjectMetadata.from_dict(object_metadatum_dict)
228 object_metadata.append(object_metadatum)
229
230 return object_metadata
231
232 def _write_manifest_files(self, storage_provider: StorageProvider, object_metadata: list[ObjectMetadata]) -> None:
233 """
234 Writes the main manifest and its part files.
235
236 :param storage_provider: The storage provider to use for writing.
237 :param object_metadata: objects to include in manifest.
238 """
239 if not object_metadata:
240 return
241
242 base_path = self._manifest_path
243 manifest_base_path = base_path
244
245 base_path_parts = base_path.split(os.sep)
246 if DEFAULT_MANIFEST_BASE_DIR in base_path_parts:
247 manifests_index = base_path_parts.index(DEFAULT_MANIFEST_BASE_DIR)
248 if manifests_index > 0:
249 manifest_base_path = os.path.join(*base_path_parts[:manifests_index])
250 else:
251 manifest_base_path = ""
252 if base_path.startswith(os.sep):
253 manifest_base_path = os.sep + manifest_base_path
254
255 current_time = datetime.now(timezone.utc)
256 current_time_str = current_time.isoformat(timespec="seconds")
257 manifest_folderpath = os.path.join(manifest_base_path, DEFAULT_MANIFEST_BASE_DIR, current_time_str)
258 # We currently write only one part by default
259 part_sequence_number = 1
260 manifest_part_file_path = os.path.join(
261 MANIFEST_PARTS_CHILD_DIR,
262 f"{MANIFEST_PART_PREFIX}{part_sequence_number:0{SEQUENCE_PADDING}}{MANIFEST_PART_SUFFIX}",
263 )
264
265 manifest = Manifest(version="1", parts=[ManifestPartReference(path=manifest_part_file_path)])
266
267 # Write single manifest part with metadata as JSON lines (each object on a new line)
268 manifest_part_content = "\n".join(
269 [json.dumps(_metadata_to_manifest_dict(metadata)) for metadata in object_metadata]
270 )
271 storage_provider.put_object(
272 os.path.join(manifest_folderpath, manifest_part_file_path), manifest_part_content.encode("utf-8")
273 )
274
275 # Write the main manifest file
276 manifest_file_path = os.path.join(manifest_folderpath, MANIFEST_INDEX_FILENAME)
277 manifest_content = manifest.to_json()
278 storage_provider.put_object(manifest_file_path, manifest_content.encode("utf-8"))
279
[docs]
280 def list_objects(
281 self,
282 prefix: str,
283 start_after: Optional[str] = None,
284 end_at: Optional[str] = None,
285 include_directories: bool = False,
286 ) -> Iterator[ObjectMetadata]:
287 if (start_after is not None) and (end_at is not None) and not (start_after < end_at):
288 raise ValueError(f"start_after ({start_after}) must be before end_at ({end_at})!")
289
290 if prefix and not prefix.endswith("/"):
291 prefix = prefix + "/"
292
293 # Note that this is a generator, not a tuple (there's no tuple comprehension).
294 keys = (
295 key
296 for key in self._files
297 if key.startswith(prefix)
298 and (start_after is None or start_after < key)
299 and (end_at is None or key <= end_at)
300 )
301
302 pending_directory: Optional[ObjectMetadata] = None
303 for key in sorted(keys):
304 if include_directories:
305 relative = key[len(prefix) :].lstrip("/")
306 subdirectory = relative.split("/", 1)[0] if "/" in relative else None
307
308 if subdirectory:
309 directory_name = f"{prefix}{subdirectory}/"
310
311 if pending_directory and pending_directory.key != directory_name:
312 yield pending_directory
313
314 obj_metadata = self.get_object_metadata(key)
315 if not pending_directory or pending_directory.key != directory_name:
316 pending_directory = ObjectMetadata(
317 key=directory_name,
318 type="directory",
319 last_modified=obj_metadata.last_modified,
320 content_length=0,
321 )
322 else:
323 pending_directory.last_modified = max(
324 pending_directory.last_modified, obj_metadata.last_modified
325 )
326 continue # Skip yielding this key as it's part of a directory
327
328 obj = self._files[key]
329 obj.key = key # use key without base_path
330 yield obj
331
332 if include_directories and pending_directory:
333 yield pending_directory
334
[docs]
335 def get_object_metadata(self, path: str, include_pending: bool = False) -> ObjectMetadata:
336 if path in self._files:
337 if include_pending and path in self._pending_removes:
338 raise FileNotFoundError(f"Object {path} does not exist.")
339 else:
340 return self._files[path]
341 elif include_pending and path in self._pending_adds:
342 return self._pending_adds[path]
343 else:
344 raise FileNotFoundError(f"Object {path} does not exist.")
345
[docs]
346 def glob(self, pattern: str) -> list[str]:
347 all_objects = [object.key for object in self.list_objects("")]
348 return [key for key in glob(all_objects, pattern)]
349
[docs]
350 def realpath(self, path: str) -> tuple[str, bool]:
351 exists = path in self._files
352 return path, exists
353
[docs]
354 def add_file(self, path: str, metadata: ObjectMetadata) -> None:
355 if not self.is_writable():
356 raise RuntimeError(f"Manifest update support not enabled in configuration. Attempted to add {path}.")
357 self._pending_adds[path] = metadata
358
[docs]
359 def remove_file(self, path: str) -> None:
360 if not self.is_writable():
361 raise RuntimeError(f"Manifest update support not enabled in configuration. Attempted to remove {path}.")
362 if path not in self._files:
363 raise FileNotFoundError(f"Object {path} does not exist.")
364 self._pending_removes.add(path)
365
368
[docs]
369 def commit_updates(self) -> None:
370 if not self._pending_adds and not self._pending_removes:
371 return
372
373 if self._pending_adds:
374 self._files.update(self._pending_adds)
375 self._pending_adds = {}
376
377 for path in self._pending_removes:
378 self._files.pop(path)
379 self._pending_removes = set()
380
381 # Collect metadata for each object to write out in this part file.
382 object_metadata = [
383 ObjectMetadata(key=file_path, content_length=metadata.content_length, last_modified=metadata.last_modified)
384 for file_path, metadata in self._files.items()
385 ]
386 self._write_manifest_files(self._storage_provider, object_metadata)