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 json
19import logging
20import os
21from collections.abc import Iterator
22from dataclasses import asdict, dataclass
23from datetime import datetime, timezone
24from typing import Any, Optional, Union
25
26from ..types import (
27 AWARE_DATETIME_MIN,
28 MAX_SYMLINK_DEPTH,
29 MetadataProvider,
30 ObjectMetadata,
31 ResolvedPath,
32 ResolvedPathState,
33 StorageProvider,
34)
35from ..utils import create_attribute_filter_evaluator, glob, matches_attribute_filter_expression
36from .manifest_formats import ManifestFormat, get_format_handler
37from .manifest_object_metadata import ManifestObjectMetadata
38
39logger = logging.getLogger(__name__)
40
41
42DEFAULT_MANIFEST_BASE_DIR = ".msc_manifests"
43MANIFEST_INDEX_FILENAME = "msc_manifest_index.json"
44MANIFEST_PARTS_CHILD_DIR = "parts"
45MANIFEST_PART_PREFIX = "msc_manifest_part"
46SEQUENCE_PADDING = 6 # Define padding for the sequence number (e.g., 6 for "000001")
47
48
[docs]
49@dataclass
50class ManifestPartReference:
51 """
52 A data class representing a reference to dataset manifest part.
53 """
54
55 #: The path of the manifest part relative to the main manifest.
56 path: str
57
[docs]
58 @staticmethod
59 def from_dict(data: dict[str, Any]) -> ManifestPartReference:
60 """
61 Creates a ManifestPartReference instance from a dictionary.
62 """
63 # Validate that the required 'path' field is present
64 if "path" not in data:
65 raise ValueError("Missing required field: 'path'")
66
67 return ManifestPartReference(path=data["path"])
68
[docs]
69 def to_dict(self) -> dict:
70 """
71 Converts ManifestPartReference instance to a dictionary.
72 """
73 return {
74 "path": self.path,
75 }
76
77
[docs]
78@dataclass
79class Manifest:
80 """
81 A data class representing a dataset manifest.
82 """
83
84 #: Defines the version of the manifest schema.
85 version: str
86 #: References to manifest parts.
87 parts: list[ManifestPartReference]
88 #: Format of manifest parts (jsonl or parquet).
89 format: str = "jsonl"
90
[docs]
91 @staticmethod
92 def from_dict(data: dict) -> "Manifest":
93 """
94 Creates a Manifest instance from a dictionary (parsed from JSON).
95 """
96 try:
97 version = data["version"]
98 parts = [ManifestPartReference.from_dict(part) for part in data["parts"]]
99 format = data.get("format", "jsonl")
100 except KeyError as e:
101 raise ValueError("Invalid manifest data: Missing required field") from e
102
103 return Manifest(version=version, parts=parts, format=format)
104
[docs]
105 def to_json(self) -> str:
106 data = asdict(self)
107 data["parts"] = [part.to_dict() for part in self.parts]
108 return json.dumps(data)
109
110
[docs]
111class ManifestMetadataProvider(MetadataProvider):
112 _storage_provider: StorageProvider
113 _files: dict[str, ManifestObjectMetadata]
114 _pending_adds: dict[str, ManifestObjectMetadata]
115 _pending_removes: set[str]
116 _manifest_path: str
117 _writable: bool
118 _allow_overwrites: bool
119 _format: Union[ManifestFormat, str]
120
121 def __init__(
122 self,
123 storage_provider: StorageProvider,
124 manifest_path: str,
125 writable: bool = False,
126 allow_overwrites: bool = False,
127 manifest_format: Union[ManifestFormat, str] = ManifestFormat.JSONL,
128 ) -> None:
129 """
130 Creates a :py:class:`ManifestMetadataProvider`.
131
132 :param storage_provider: Storage provider.
133 :param manifest_path: Main manifest file path.
134 :param writable: If true, allows modifications and new manifests to be written.
135 :param allow_overwrites: If true, allows overwriting existing files without error.
136 :param manifest_format: Format for manifest parts. Defaults to ManifestFormat.JSONL.
137 """
138 self._storage_provider = storage_provider
139 self._files = {}
140 self._pending_adds = {}
141 self._pending_removes = set()
142 self._manifest_path = manifest_path
143 self._writable = writable
144 self._allow_overwrites = allow_overwrites
145 self._format = (
146 manifest_format if isinstance(manifest_format, ManifestFormat) else ManifestFormat(manifest_format)
147 )
148
149 self._load_manifest(storage_provider, self._manifest_path)
150
151 def _load_manifest(self, storage_provider: StorageProvider, manifest_path: str) -> None:
152 """
153 Loads manifest.
154
155 :param storage_provider: Storage provider.
156 :param manifest_path: Main manifest file path
157 """
158
159 def helper_find_manifest_file(manifest_path: str) -> str:
160 if storage_provider.is_file(manifest_path):
161 return manifest_path
162
163 if storage_provider.is_file(os.path.join(manifest_path, MANIFEST_INDEX_FILENAME)):
164 return os.path.join(manifest_path, MANIFEST_INDEX_FILENAME)
165
166 # Now go looking and select newest manifest.
167 if DEFAULT_MANIFEST_BASE_DIR not in manifest_path.split("/"):
168 manifest_path = os.path.join(manifest_path, DEFAULT_MANIFEST_BASE_DIR)
169
170 candidates = storage_provider.glob(os.path.join(manifest_path, "*", MANIFEST_INDEX_FILENAME))
171 candidates = sorted(candidates)
172 return candidates[-1] if candidates else ""
173
174 resolved_manifest_path = helper_find_manifest_file(manifest_path)
175 if not resolved_manifest_path:
176 logger.warning(f"No manifest found at '{manifest_path}'.")
177 return
178
179 file_content = storage_provider.get_object(resolved_manifest_path)
180
181 prefix = os.path.dirname(resolved_manifest_path)
182 _, file_extension = os.path.splitext(resolved_manifest_path)
183 self._load_manifest_file(storage_provider, file_content, prefix, file_extension[1:])
184
185 def _load_manifest_file(
186 self, storage_provider: StorageProvider, file_content: bytes, manifest_base: str, file_type: str
187 ) -> None:
188 """
189 Loads a manifest.
190
191 :param storage_provider: Storage provider.
192 :param file_content: Manifest file content bytes.
193 :param manifest_base: Manifest file base path.
194 :param file_type: Manifest file type.
195 """
196 if file_type == "json":
197 manifest_dict = json.loads(file_content.decode("utf-8"))
198 manifest = Manifest.from_dict(manifest_dict)
199
200 # Check manifest version. Not needed once we make the manifest model use sum types/discriminated unions.
201 if manifest.version != "1":
202 raise ValueError(f"Manifest version {manifest.version} is not supported.")
203
204 # Load manifest parts.
205 for manifest_part_reference in manifest.parts:
206 object_metadata_list: list[ManifestObjectMetadata] = self._load_manifest_part_file(
207 storage_provider=storage_provider,
208 manifest_base=manifest_base,
209 manifest_part_reference=manifest_part_reference,
210 manifest_format=manifest.format,
211 )
212
213 for object_metadatum in object_metadata_list:
214 self._files[object_metadatum.key] = object_metadatum
215 else:
216 raise NotImplementedError(f"Manifest file type {file_type} is not supported.")
217
218 def _load_manifest_part_file(
219 self,
220 storage_provider: StorageProvider,
221 manifest_base: str,
222 manifest_part_reference: ManifestPartReference,
223 manifest_format: Union[ManifestFormat, str] = ManifestFormat.JSONL,
224 ) -> list[ManifestObjectMetadata]:
225 """
226 Loads a manifest part and converts to ManifestObjectMetadata.
227
228 :param storage_provider: Storage provider.
229 :param manifest_base: Manifest file base path. Prepend to manifest part reference paths.
230 :param manifest_part_reference: Manifest part reference.
231 :param manifest_format: Format of the manifest part (jsonl or parquet).
232 """
233 if not os.path.isabs(manifest_part_reference.path):
234 remote_path = os.path.join(manifest_base, manifest_part_reference.path)
235 else:
236 remote_path = manifest_part_reference.path
237 manifest_part_file_content = storage_provider.get_object(remote_path)
238
239 _, ext = os.path.splitext(remote_path)
240 detected_format = ext[1:] if ext else manifest_format
241
242 format_handler = get_format_handler(detected_format)
243 object_metadata_list = format_handler.read_part(manifest_part_file_content)
244
245 # Convert ObjectMetadata to ManifestObjectMetadata
246 # The format handler returns ObjectMetadata, but may have extra attributes set (like physical_path)
247 return [ManifestObjectMetadata.from_object_metadata(obj) for obj in object_metadata_list]
248
249 def _write_manifest_files(
250 self,
251 storage_provider: StorageProvider,
252 object_metadata: list[ManifestObjectMetadata],
253 manifest_format: Union[ManifestFormat, str] = ManifestFormat.JSONL,
254 ) -> None:
255 """
256 Writes the main manifest and its part files.
257
258 Accepts ManifestObjectMetadata which extends ObjectMetadata, so format handlers
259 (which expect ObjectMetadata) can serialize it, preserving all fields including physical_path.
260
261 :param storage_provider: The storage provider to use for writing.
262 :param object_metadata: ManifestObjectMetadata objects to include in manifest.
263 :param manifest_format: Format for manifest parts. Defaults to ManifestFormat.JSONL.
264 """
265 if not object_metadata:
266 return
267
268 base_path = self._manifest_path
269 manifest_base_path = base_path
270
271 base_path_parts = base_path.split(os.sep)
272 if DEFAULT_MANIFEST_BASE_DIR in base_path_parts:
273 manifests_index = base_path_parts.index(DEFAULT_MANIFEST_BASE_DIR)
274 if manifests_index > 0:
275 manifest_base_path = os.path.join(*base_path_parts[:manifests_index])
276 else:
277 manifest_base_path = ""
278 if base_path.startswith(os.sep):
279 manifest_base_path = os.sep + manifest_base_path
280
281 current_time = datetime.now(timezone.utc)
282 current_time_str = current_time.isoformat(timespec="seconds")
283 manifest_folderpath = os.path.join(manifest_base_path, DEFAULT_MANIFEST_BASE_DIR, current_time_str)
284
285 format_handler = get_format_handler(manifest_format)
286 suffix = format_handler.get_file_suffix()
287
288 # We currently write only one part by default.
289 part_sequence_number = 1
290 manifest_part_file_path = os.path.join(
291 MANIFEST_PARTS_CHILD_DIR,
292 f"{MANIFEST_PART_PREFIX}{part_sequence_number:0{SEQUENCE_PADDING}}{suffix}",
293 )
294
295 format_value = manifest_format.value if isinstance(manifest_format, ManifestFormat) else manifest_format
296 manifest = Manifest(
297 version="1", parts=[ManifestPartReference(path=manifest_part_file_path)], format=format_value
298 )
299
300 manifest_part_content = format_handler.write_part(list(object_metadata))
301 storage_provider.put_object(os.path.join(manifest_folderpath, manifest_part_file_path), manifest_part_content)
302
303 manifest_file_path = os.path.join(manifest_folderpath, MANIFEST_INDEX_FILENAME)
304 manifest_content = manifest.to_json()
305 storage_provider.put_object(manifest_file_path, manifest_content.encode("utf-8"))
306
[docs]
307 def list_objects(
308 self,
309 path: str,
310 start_after: Optional[str] = None,
311 end_at: Optional[str] = None,
312 include_directories: bool = False,
313 attribute_filter_expression: Optional[str] = None,
314 show_attributes: bool = False,
315 ) -> Iterator[ObjectMetadata]:
316 """
317 List objects in the manifest.
318
319 :param path: The path to filter objects by.
320 :param start_after: The object to start after.
321 :param end_at: The object to end at.
322 :param include_directories: Whether to include directories.
323 :param attribute_filter_expression: The attribute filter expression to filter objects by.
324 :param show_attributes: This field is not used in this implementation - It will always return attributes. This is present merely to satisfy the interface.
325 """
326
327 if (start_after is not None) and (end_at is not None) and not (start_after < end_at):
328 raise ValueError(f"start_after ({start_after}) must be before end_at ({end_at})!")
329
330 if path and not path.endswith("/"):
331 path = path + "/"
332
333 # create evaluator for attribute filter expression if present
334 evaluator = (
335 create_attribute_filter_evaluator(attribute_filter_expression) if attribute_filter_expression else None
336 )
337
338 if include_directories:
339 # Match delimiter-style providers such as S3: list the immediate child
340 # directories first, then the immediate files. The manifest only stores
341 # object keys, so this view is built on demand for the requested path.
342 directories: dict[str, ObjectMetadata] = {}
343 synthetic_directories: set[str] = set()
344 files: list[ObjectMetadata] = []
345
346 for key, obj_metadata in sorted(self._files.items()):
347 if not key.startswith(path) or (
348 evaluator is not None and not matches_attribute_filter_expression(obj_metadata, evaluator)
349 ):
350 continue
351
352 relative = key[len(path) :].lstrip("/")
353
354 if "/" not in relative:
355 if (start_after is not None and key <= start_after) or (end_at is not None and end_at < key):
356 continue
357
358 if obj_metadata.type == "directory":
359 # Prefer explicit directory entries from the manifest when present.
360 directory_key = obj_metadata.key.rstrip("/")
361 directories[directory_key] = ObjectMetadata(
362 key=directory_key,
363 type="directory",
364 last_modified=obj_metadata.last_modified,
365 content_length=0,
366 content_type=obj_metadata.content_type,
367 etag=obj_metadata.etag,
368 storage_class=obj_metadata.storage_class,
369 metadata=obj_metadata.metadata,
370 symlink_target=obj_metadata.symlink_target,
371 )
372 synthetic_directories.discard(directory_key)
373 else:
374 files.append(obj_metadata)
375 continue
376
377 # Nested files imply an immediate child directory under the listed path.
378 directory_name = f"{path}{relative.split('/', 1)[0]}"
379 directory_key = directory_name.rstrip("/")
380 existing_directory = directories.get(directory_key)
381 if existing_directory is None:
382 directories[directory_key] = ObjectMetadata(
383 key=directory_key,
384 type="directory",
385 last_modified=obj_metadata.last_modified,
386 content_length=0,
387 )
388 synthetic_directories.add(directory_key)
389 else:
390 if directory_key in synthetic_directories:
391 existing_directory.last_modified = max(
392 existing_directory.last_modified, obj_metadata.last_modified
393 )
394
395 yield from (
396 directories[key]
397 for key in sorted(directories)
398 if (start_after is None or start_after < key) and (end_at is None or key <= end_at)
399 )
400 yield from sorted(files, key=lambda obj: obj.key)
401 return
402
403 # Note that this is a generator, not a tuple (there's no tuple comprehension).
404 keys = (
405 key
406 for key, obj_metadata in self._files.items()
407 if key.startswith(path)
408 and obj_metadata.type != "directory"
409 and (start_after is None or start_after < key)
410 and (end_at is None or key <= end_at)
411 and (
412 evaluator is None or matches_attribute_filter_expression(obj_metadata, evaluator)
413 ) # filter by evaluator if present
414 )
415
416 for key in sorted(keys):
417 yield self._files[key]
418
[docs]
419 def get_object_metadata(self, path: str, include_pending: bool = False) -> ObjectMetadata:
420 if path in self._files:
421 if include_pending and path in self._pending_removes:
422 raise FileNotFoundError(f"Object {path} does not exist.")
423 else:
424 # Return ManifestObjectMetadata directly (it extends ObjectMetadata)
425 return self._files[path]
426 elif include_pending and path in self._pending_adds:
427 return self._pending_adds[path]
428 else:
429 return self._get_directory_metadata(path, include_pending)
430
431 def _get_directory_metadata(self, path: str, include_pending: bool) -> ObjectMetadata:
432 dir_prefix = path.rstrip("/") + "/"
433
434 has_committed = any(
435 k.startswith(dir_prefix) and (not include_pending or k not in self._pending_removes) for k in self._files
436 )
437 has_pending = include_pending and any(k.startswith(dir_prefix) for k in self._pending_adds)
438
439 if not has_committed and not has_pending:
440 raise FileNotFoundError(f"Object {path} does not exist.")
441
442 return ObjectMetadata(
443 key=dir_prefix,
444 type="directory",
445 content_length=0,
446 last_modified=AWARE_DATETIME_MIN,
447 )
448
[docs]
449 def glob(self, pattern: str, attribute_filter_expression: Optional[str] = None) -> list[str]:
450 """
451 List objects in the manifest.
452
453 :param pattern: The pattern to filter objects by.
454 :param attribute_filter_expression: The attribute filter expression to filter objects by.
455 """
456
457 all_objects = [
458 object.key for object in self.list_objects("", attribute_filter_expression=attribute_filter_expression)
459 ]
460 return [key for key in glob(all_objects, pattern)]
461
[docs]
462 def realpath(self, logical_path: str) -> ResolvedPath:
463 """
464 Resolves a logical path to its physical storage path if the object exists.
465 Only checks committed files, not pending changes.
466 Follows ``symlink_target`` chains in the manifest with depth limit and cycle detection.
467
468 :param logical_path: The user-facing logical path
469
470 :return: ResolvedPath with exists=True if found, exists=False otherwise
471 """
472 visited: set[str] = set()
473 current = logical_path
474 for _ in range(MAX_SYMLINK_DEPTH):
475 manifest_obj = self._files.get(current)
476 if not manifest_obj:
477 return ResolvedPath(physical_path=current, state=ResolvedPathState.UNTRACKED, profile=None)
478 if not manifest_obj.symlink_target:
479 if manifest_obj.physical_path is None:
480 raise ValueError(f"Manifest entry for '{current}' has no physical path")
481 return ResolvedPath(
482 physical_path=manifest_obj.physical_path, state=ResolvedPathState.EXISTS, profile=None
483 )
484 if current in visited:
485 raise ValueError(f"Symlink cycle detected at: {current}")
486 visited.add(current)
487 current = ObjectMetadata.resolve_symlink_target(current, manifest_obj.symlink_target)
488 raise ValueError(f"Too many levels of symlinks (>{MAX_SYMLINK_DEPTH}): {logical_path}")
489
[docs]
490 def generate_physical_path(self, logical_path: str, for_overwrite: bool = False) -> ResolvedPath:
491 """
492 Generates a physical storage path for a new object or for overwriting an existing object.
493
494 For now, this simply returns the logical path (no path rewriting).
495 In the future, this could generate unique paths for overwrites.
496
497 :param logical_path: The user-facing logical path
498 :param for_overwrite: When ``True``, generate a path for overwriting an existing object.
499
500 :return: The physical storage path to use for writing
501 """
502 # For now, physical path = logical path
503 # Future enhancement: generate unique paths for overwrites
504 # if for_overwrite and self._allow_overwrites:
505 # return f"{logical_path}-{uuid.uuid4().hex}"
506 return ResolvedPath(physical_path=logical_path, state=ResolvedPathState.UNTRACKED, profile=None)
507
[docs]
508 def add_file(self, path: str, metadata: ObjectMetadata) -> None:
509 if not self.is_writable():
510 raise RuntimeError(f"Manifest update support not enabled in configuration. Attempted to add {path}.")
511
512 # Check if file already exists in committed state and overwrites are not allowed
513 # Pending files can always be overwritten
514 if not self._allow_overwrites and path in self._files:
515 raise FileExistsError(f"File {path} already exists and overwrites are not allowed.")
516
517 # Handle two cases:
518 # 1. If metadata is already a ManifestObjectMetadata, use it directly
519 # 2. Otherwise, create one assuming metadata.key contains the physical path
520 if isinstance(metadata, ManifestObjectMetadata):
521 # Already has proper logical/physical path separation
522 manifest_metadata = metadata
523 # Ensure the logical path matches what was requested
524 if manifest_metadata.key != path:
525 raise ValueError(f"Logical path mismatch: expected {path}, got {manifest_metadata.key}")
526 else:
527 # For backward compatibility, create a ManifestObjectMetadata from ObjectMetadata
528 manifest_metadata = ManifestObjectMetadata(
529 key=path, # Logical path (user-facing)
530 content_length=metadata.content_length,
531 last_modified=metadata.last_modified,
532 content_type=metadata.content_type,
533 etag=metadata.etag,
534 metadata=metadata.metadata,
535 type=metadata.type,
536 physical_path=metadata.key, # Physical path (from storage provider)
537 )
538
539 # TODO: Time travel is not supported - we do not rename file paths when overwriting.
540 # When a file is overwritten, the manifest points to the new version of the file at the same location
541 # without maintaining history of previous versions.
542 self._pending_adds[path] = manifest_metadata
543
[docs]
544 def remove_file(self, path: str) -> None:
545 if not self.is_writable():
546 raise RuntimeError(f"Manifest update support not enabled in configuration. Attempted to remove {path}.")
547 if path not in self._files:
548 raise FileNotFoundError(f"Object {path} does not exist.")
549 self._pending_removes.add(path)
550
553
556
559
[docs]
560 def commit_updates(self) -> None:
561 if not self._pending_adds and not self._pending_removes:
562 return
563
564 if self._pending_adds:
565 self._files.update(self._pending_adds)
566 self._pending_adds = {}
567
568 for path in self._pending_removes:
569 self._files.pop(path)
570 self._pending_removes = set()
571
572 # Serialize ManifestObjectMetadata directly
573 # to_dict() will include all fields including physical_path
574 object_metadata = list(self._files.values())
575 self._write_manifest_files(self._storage_provider, object_metadata, manifest_format=self._format)