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
17
18import posixpath
19from abc import ABC, abstractmethod
20from collections.abc import Iterator, Sequence
21from dataclasses import asdict, dataclass, field, replace
22from datetime import datetime, timezone
23from enum import Enum
24from typing import IO, Any, NamedTuple
25
26from dateutil.parser import parse as dateutil_parser
27
28MSC_PROTOCOL_NAME = "msc"
29MSC_PROTOCOL = MSC_PROTOCOL_NAME + "://"
30
31DEFAULT_RETRY_ATTEMPTS = 3
32DEFAULT_RETRY_DELAY = 1.0
33DEFAULT_RETRY_BACKOFF_MULTIPLIER = 2.0
34
35# datetime.min is a naive datetime.
36#
37# This creates issues when doing datetime.astimezone(timezone.utc) since it assumes the local timezone for the naive datetime.
38# If the local timezone is offset behind UTC, it attempts to subtract off the offset which goes below the representable limit (i.e. an underflow).
39# A `ValueError: year 0 is out of range` is thrown as a result.
40AWARE_DATETIME_MIN = datetime.min.replace(tzinfo=timezone.utc)
41
42
[docs]
43class SymlinkHandling(str, Enum):
44 """Controls how symbolic links are handled during listing and sync operations.
45
46 External symlinks (targets outside ``base_path``) are skipped by default in
47 ``FOLLOW`` and ``PRESERVE``. Use ``FOLLOW_STRICT`` or ``PRESERVE_STRICT`` to
48 fail instead.
49
50 - ``FOLLOW``: Dereference internal symlinks; skip external symlinks. This is
51 the default.
52 - ``FOLLOW_STRICT``: Dereference internal symlinks; raise on external symlinks.
53 - ``SKIP``: All symlinks are excluded from results.
54 - ``PRESERVE``: Surface internal symlinks as leaf entries with
55 :py:attr:`ObjectMetadata.symlink_target` populated; skip external symlinks.
56 Directory symlinks are **not** recursed into.
57 - ``PRESERVE_STRICT``: Same as ``PRESERVE`` for internal symlinks; raise on
58 external symlinks.
59
60 .. note::
61 This option is only meaningful for POSIX file storage providers, which
62 have native symlink semantics. Cloud storage providers (S3, GCS, Azure,
63 OCI, AIS) ignore this parameter: they always list whatever is in the
64 bucket/prefix and surface MSC's symlink convention (an empty object
65 whose user metadata carries the target path) with
66 :py:attr:`ObjectMetadata.symlink_target` populated -- object storage
67 isn't a filesystem, so there is nothing to "follow" during listing.
68 """
69
70 FOLLOW = "follow"
71 FOLLOW_STRICT = "follow_strict"
72 SKIP = "skip"
73 PRESERVE = "preserve"
74 PRESERVE_STRICT = "preserve_strict"
75
76
77# Maximum number of symlink hops allowed when resolving a symlink chain.
78MAX_SYMLINK_DEPTH = 8
79
80
[docs]
81class SignerType(str, Enum):
82 """Supported signer backends for presigned URL generation."""
83
84 S3 = "s3"
85 CLOUDFRONT = "cloudfront"
86 AZURE = "azure"
87
88
[docs]
89@dataclass
90class Credentials:
91 """
92 A data class representing the credentials needed to access a storage provider.
93 """
94
95 #: The access key for authentication.
96 access_key: str
97 #: The secret key for authentication.
98 secret_key: str
99 #: An optional security token for temporary credentials.
100 token: str | None
101 #: The expiration time of the credentials in ISO 8601 format.
102 expiration: str | None
103 #: A dictionary for storing custom key-value pairs.
104 custom_fields: dict[str, Any] = field(default_factory=dict)
105
[docs]
106 def is_expired(self) -> bool:
107 """
108 Checks if the credentials are expired based on the expiration time.
109
110 :return: ``True`` if the credentials are expired, ``False`` otherwise.
111 """
112 expiry = dateutil_parser(self.expiration) if self.expiration else None
113 if expiry is None:
114 return False
115 return expiry <= datetime.now(tz=timezone.utc)
116
[docs]
117 def get_custom_field(self, key: str, default: Any = None) -> Any:
118 """
119 Retrieves a value from custom fields by its key.
120
121 :param key: The key to look up in custom fields.
122 :param default: The default value to return if the key is not found.
123 :return: The value associated with the key, or the default value if not found.
124 """
125 return self.custom_fields.get(key, default)
126
127
227
228
[docs]
229class CredentialsProvider(ABC):
230 """
231 Abstract base class for providing credentials to access a storage provider.
232 """
233
[docs]
234 @abstractmethod
235 def get_credentials(self) -> Credentials:
236 """
237 Retrieves the current credentials.
238
239 :return: The current credentials used for authentication.
240 """
241
[docs]
242 @abstractmethod
243 def refresh_credentials(self) -> None:
244 """
245 Refreshes the credentials if they are expired or about to expire.
246 """
247
248
[docs]
249@dataclass
250class Range:
251 """
252 A data class that represents a byte range for read operations.
253 """
254
255 #: The start offset in bytes.
256 offset: int
257 #: The number of bytes to read.
258 size: int
259
260
[docs]
261class StorageProvider(ABC):
262 """
263 Abstract base class for interacting with a storage provider.
264 """
265
[docs]
266 @abstractmethod
267 def put_object(
268 self,
269 path: str,
270 body: bytes,
271 if_match: str | None = None,
272 if_none_match: str | None = None,
273 attributes: dict[str, Any] | None = None,
274 ) -> None:
275 """
276 Uploads an object to the storage provider.
277
278 :param path: The path where the object will be stored.
279 :param body: The content of the object to store.
280 :param if_match: Optional If-Match value for conditional upload.
281 :param if_none_match: Optional If-None-Match value for conditional upload.
282 :param attributes: The attributes to add to the file.
283 """
284
[docs]
285 @abstractmethod
286 def get_object(self, path: str, byte_range: Range | None = None) -> bytes:
287 """
288 Retrieves an object from the storage provider.
289
290 :param path: The path where the object is stored.
291 :param byte_range: Optional byte range (offset, length) to read.
292 :return: The content of the retrieved object.
293 """
294
[docs]
295 @abstractmethod
296 def copy_object(self, src_path: str, dest_path: str) -> None:
297 """
298 Copies an object from source to destination in the storage provider.
299
300 :param src_path: The path of the source object to copy.
301 :param dest_path: The path of the destination.
302 """
303
[docs]
304 @abstractmethod
305 def delete_object(self, path: str, if_match: str | None = None) -> None:
306 """
307 Deletes an object from the storage provider.
308
309 :param path: The path of the object to delete.
310 :param if_match: Optional if-match value to use for conditional deletion.
311 """
312
[docs]
313 @abstractmethod
314 def delete_objects(self, paths: list[str]) -> None:
315 """
316 Deletes multiple objects from the storage provider.
317
318 :param paths: A list of paths of objects to delete.
319 """
320
[docs]
321 @abstractmethod
322 def make_symlink(self, path: str, target: str) -> None:
323 """
324 Creates a symbolic link at ``path`` pointing to ``target``.
325
326 On POSIX backends this creates a native OS symlink with a relative target path.
327 On object-store backends this creates a zero-byte marker object with the target
328 stored in user metadata (``msc-symlink-target``).
329
330 :param path: The path where the symlink will be created.
331 :param target: The logical key that the symlink points to.
332 """
333
344
[docs]
345 @abstractmethod
346 def list_objects(
347 self,
348 path: str,
349 start_after: str | None = None,
350 end_at: str | None = None,
351 include_directories: bool = False,
352 attribute_filter_expression: str | None = None,
353 show_attributes: bool = False,
354 symlink_handling: SymlinkHandling = SymlinkHandling.FOLLOW,
355 ) -> Iterator[ObjectMetadata]:
356 """
357 Lists objects in the storage provider under the specified path.
358
359 :param path: The path to list objects under. The path must be a valid file or subdirectory path, cannot be partial or just "prefix".
360 :param start_after: The key to start after (i.e. exclusive). An object with this key doesn't have to exist.
361 :param end_at: The key to end at (i.e. inclusive). An object with this key doesn't have to exist.
362 :param include_directories: Whether to include directories in the result. When ``True``, directories are returned alongside objects.
363 :param attribute_filter_expression: The attribute filter expression to apply to the result.
364 :param show_attributes: Whether to return attributes in the result. There will be performance impact if this is True as now we need to get object metadata for each object.
365 :param symlink_handling: How to handle symbolic links during listing.
366
367 :return: An iterator over objects metadata under the specified path.
368 """
369
[docs]
370 @abstractmethod
371 def list_objects_recursive(
372 self,
373 path: str = "",
374 start_after: str | None = None,
375 end_at: str | None = None,
376 max_workers: int = 32,
377 look_ahead: int = 2,
378 symlink_handling: SymlinkHandling = SymlinkHandling.FOLLOW,
379 ) -> Iterator[ObjectMetadata]:
380 """
381 Lists files recursively in the storage provider under the specified path.
382
383 :param path: The path to list objects under.
384 :param start_after: The key to start after (i.e. exclusive). An object with this key doesn't have to exist.
385 :param end_at: The key to end at (i.e. inclusive). An object with this key doesn't have to exist.
386 :param max_workers: Maximum concurrent workers for provider-level recursive listing.
387 :param look_ahead: Prefixes to buffer per worker for provider-level recursive listing.
388 :param symlink_handling: How to handle symbolic links during listing.
389 :return: An iterator over object metadata under the specified path.
390 """
391
[docs]
392 @abstractmethod
393 def upload_file(self, remote_path: str, f: str | IO, attributes: dict[str, Any] | None = None) -> None:
394 """
395 Uploads a file from the local file system to the storage provider.
396
397 :param remote_path: The path where the object will be stored.
398 :param f: The source file to upload. This can either be a string representing the local
399 file path, or a file-like object (e.g., an open file handle).
400 :param attributes: The attributes to add to the file if a new file is created.
401 """
402
[docs]
403 @abstractmethod
404 def download_file(self, remote_path: str, f: str | IO, metadata: ObjectMetadata | None = None) -> None:
405 """
406 Downloads a file from the storage provider to the local file system.
407
408 :param remote_path: The path of the file to download.
409 :param f: The destination for the downloaded file. This can either be a string representing
410 the local file path where the file will be saved, or a file-like object to write the
411 downloaded content into.
412 :param metadata: Metadata about the object to download.
413 """
414
[docs]
415 @abstractmethod
416 def download_files(
417 self,
418 remote_paths: list[str],
419 local_paths: list[str],
420 metadata: Sequence[ObjectMetadata | None] | None = None,
421 max_workers: int = 16,
422 ) -> None:
423 """
424 Downloads multiple files from the storage provider to the local file system.
425
426 :param remote_paths: List of remote paths of files to download.
427 :param local_paths: List of local file paths to save the downloaded files to.
428 :param metadata: Optional per-file metadata used to decide between regular and multipart download.
429 :param max_workers: Maximum number of concurrent download workers (default: 16).
430 :raises ValueError: If remote_paths and local_paths have different lengths.
431 """
432
[docs]
433 @abstractmethod
434 def upload_files(
435 self,
436 local_paths: list[str],
437 remote_paths: list[str],
438 attributes: Sequence[dict[str, Any] | None] | None = None,
439 max_workers: int = 16,
440 ) -> None:
441 """
442 Uploads multiple files from the local file system to the storage provider.
443
444 :param local_paths: List of local file paths to upload.
445 :param remote_paths: List of remote paths to upload the files to.
446 :param attributes: Optional list of per-file attributes to add. When provided, must have the same length
447 as local_paths/remote_paths. Each element may be ``None`` for files that need no attributes.
448 :param max_workers: Maximum number of concurrent upload workers (default: 16).
449 :raises ValueError: If local_paths and remote_paths have different lengths.
450 :raises ValueError: If attributes is provided and has a different length than local_paths.
451 """
452
[docs]
453 @abstractmethod
454 def glob(self, pattern: str, attribute_filter_expression: str | None = None) -> list[str]:
455 """
456 Matches and retrieves a list of object keys in the storage provider that match the specified pattern.
457
458 :param pattern: The pattern to match object keys against, supporting wildcards (e.g., ``*.txt``).
459 :param attribute_filter_expression: The attribute filter expression to apply to the result.
460
461 :return: A list of object keys that match the specified pattern.
462 """
463
[docs]
464 @abstractmethod
465 def is_file(self, path: str) -> bool:
466 """
467 Checks whether the specified key in the storage provider points to a file (as opposed to a folder or directory).
468
469 :param path: The path to check.
470
471 :return: ``True`` if the key points to a file, ``False`` if it points to a directory or folder.
472 """
473
[docs]
474 def generate_presigned_url(
475 self,
476 path: str,
477 *,
478 method: str = "GET",
479 signer_type: SignerType | None = None,
480 signer_options: dict[str, Any] | None = None,
481 ) -> str:
482 """
483 Generate a pre-signed URL granting temporary access to the object at *path*.
484
485 :param path: The object path within the storage provider.
486 :param method: The HTTP method the URL should authorise (e.g. ``"GET"``, ``"PUT"``).
487 :param signer_type: The signing backend to use. ``None`` means the provider's native signer.
488 :param signer_options: Backend-specific options forwarded to the signer.
489 :return: A pre-signed URL string.
490 :raises NotImplementedError: If this storage provider does not support presigned URLs.
491 """
492 raise NotImplementedError(f"{type(self).__name__} does not support presigned URL generation.")
493
494
[docs]
495class ResolvedPathState(str, Enum):
496 """
497 Enum representing the state of a resolved path.
498 """
499
500 EXISTS = "exists" # File currently exists
501 DELETED = "deleted" # File existed before but has been deleted
502 UNTRACKED = "untracked" # File never existed or was never tracked
503
504
[docs]
505class ResolvedPath(NamedTuple):
506 """
507 Result of resolving a virtual path to a physical path.
508
509 :param physical_path: The physical path in storage backend
510 :param state: The state of the path (EXISTS, DELETED, or UNTRACKED)
511 :param profile: Optional profile name for routing in CompositeStorageClient.
512 None means use current client's storage provider.
513 String means route to named child StorageClient.
514
515 State meanings:
516 - EXISTS: File currently exists in metadata
517 - DELETED: File existed before but has been deleted (soft delete)
518 - UNTRACKED: File never existed or was never tracked
519 """
520
521 physical_path: str
522 state: ResolvedPathState
523 profile: str | None = None
524
525 @property
526 def exists(self) -> bool:
527 """Backward compatibility property: True if state is EXISTS."""
528 return self.state == ResolvedPathState.EXISTS
529
530
676
677
[docs]
678@dataclass
679class StorageProviderConfig:
680 """
681 A data class that represents the configuration needed to initialize a storage provider.
682 """
683
684 #: The name or type of the storage provider (e.g., ``s3``, ``gcs``, ``oci``, ``azure``).
685 type: str
686 #: Additional options required to configure the storage provider (e.g., endpoint URLs, region, etc.).
687 options: dict[str, Any] | None = None
688
689
[docs]
690@dataclass
691class StorageBackend:
692 """
693 Represents configuration for a single storage backend.
694 """
695
696 storage_provider_config: StorageProviderConfig
697 credentials_provider: CredentialsProvider | None = None
698 replicas: list[Replica] = field(default_factory=list)
699
700
[docs]
701class ProviderBundle(ABC):
702 """
703 Abstract base class that serves as a container for various providers (storage, credentials, and metadata)
704 that interact with a storage service. The :py:class:`ProviderBundle` abstracts access to these providers, allowing for
705 flexible implementations of cloud storage solutions.
706 """
707
708 @property
709 @abstractmethod
710 def storage_provider_config(self) -> StorageProviderConfig:
711 """
712 :return: The configuration for the storage provider, which includes the provider
713 name/type and additional options.
714 """
715
716 @property
717 @abstractmethod
718 def credentials_provider(self) -> CredentialsProvider | None:
719 """
720 :return: The credentials provider responsible for managing authentication credentials
721 required to access the storage service.
722 """
723
724 @property
725 @abstractmethod
726 def metadata_provider(self) -> MetadataProvider | None:
727 """
728 :return: The metadata provider responsible for retrieving metadata about objects in the storage service.
729 """
730
731 @property
732 @abstractmethod
733 def replicas(self) -> list[Replica]:
734 """
735 :return: The replicas configuration for this provider bundle, if any.
736 """
737
738
[docs]
739class ProviderBundleV2(ABC):
740 """
741 Abstract base class that serves as a container for various providers (storage, credentials, and metadata)
742 that interact with one or multiple storage service. The :py:class:`ProviderBundleV2` abstracts access to these providers, allowing for
743 flexible implementations of cloud storage solutions.
744
745 """
746
747 @property
748 @abstractmethod
749 def storage_backends(self) -> dict[str, StorageBackend]:
750 """
751 :return: Mapping of storage backend name -> StorageBackend. Must have at least one backend.
752 """
753
754 @property
755 @abstractmethod
756 def metadata_provider(self) -> MetadataProvider | None:
757 """
758 :return: The metadata provider responsible for retrieving metadata about objects in the storage service. If there are multiple backends, this is required.
759 """
760
761
[docs]
762@dataclass
763class RetryConfig:
764 """
765 A data class that represents the configuration for retry strategy.
766 """
767
768 #: The number of attempts before giving up. Must be at least 1.
769 attempts: int = DEFAULT_RETRY_ATTEMPTS
770 #: The base delay (in seconds) for exponential backoff. Must be a non-negative value.
771 delay: float = DEFAULT_RETRY_DELAY
772 #: The backoff multiplier for exponential backoff. Must be at least 1.0.
773 backoff_multiplier: float = DEFAULT_RETRY_BACKOFF_MULTIPLIER
774
775 def __post_init__(self) -> None:
776 if self.attempts < 1:
777 raise ValueError("Attempts must be at least 1.")
778 if self.delay < 0:
779 raise ValueError("Delay must be a non-negative number.")
780 if self.backoff_multiplier < 1.0:
781 raise ValueError("Backoff multiplier must be at least 1.0.")
782
783
[docs]
784class RetryableError(Exception):
785 """
786 Exception raised for errors that should trigger a retry.
787 """
788
789
[docs]
790@dataclass(frozen=True)
791class BatchTransferFailure:
792 """
793 A failed item from a batch upload or download operation.
794 """
795
796 #: The item index in the batch request that failed.
797 index: int
798 #: The source path for the failed transfer.
799 source_path: str
800 #: The destination path for the failed transfer.
801 destination_path: str
802 #: The underlying exception raised for this item.
803 error: Exception
804
805
[docs]
806class BatchTransferError(Exception):
807 """
808 Exception raised when one or more items fail in a batch transfer operation.
809 """
810
811 def __init__(self, failures: Sequence[BatchTransferFailure]):
812 if not failures:
813 raise ValueError("BatchTransferError requires at least one failure.")
814 self.failures = list(failures)
815 super().__init__(self._format_message())
816
817 def _format_message(self) -> str:
818 details = ", ".join(
819 f"index {failure.index} ({type(failure.error).__name__}: {failure.error})" for failure in self.failures[:3]
820 )
821 if len(self.failures) > 3:
822 details += f", and {len(self.failures) - 3} more"
823 return f"{len(self.failures)} batch transfer item(s) failed: {details}"
824
825
[docs]
826class PreconditionFailedError(Exception):
827 """
828 Exception raised when a precondition fails. e.g. if-match, if-none-match, etc.
829 """
830
831
[docs]
832class NotModifiedError(Exception):
833 """
834 Raised when a conditional operation fails because the resource has not been modified.
835
836 This typically occurs when using if-none-match with a specific generation/etag
837 and the resource's current generation/etag matches the specified one.
838 """
839
840
[docs]
841class SourceVersionCheckMode(Enum):
842 """
843 Enum for controlling source version checking behavior.
844 """
845
846 INHERIT = "inherit" # Inherit from configuration (cache config)
847 ENABLE = "enable" # Always check source version
848 DISABLE = "disable" # Never check source version
849
850
[docs]
851@dataclass
852class Replica:
853 """
854 A tier of storage that can be used to store data.
855 """
856
857 replica_profile: str
858 read_priority: int
859
860
[docs]
861class AutoCommitConfig:
862 """
863 A data class that represents the configuration for auto commit.
864 """
865
866 interval_minutes: float | None # The interval in minutes for auto commit.
867 at_exit: bool = False # if True, commit on program exit
868
869 def __init__(self, interval_minutes: float | None = None, at_exit: bool = False) -> None:
870 self.interval_minutes = interval_minutes
871 self.at_exit = at_exit
872
873
[docs]
874class ExecutionMode(Enum):
875 """
876 Enum for controlling execution mode in sync operations.
877 """
878
879 LOCAL = "local"
880 RAY = "ray"
881
882
[docs]
883class PatternType(Enum):
884 """
885 Type of pattern operation for include/exclude filtering.
886 """
887
888 INCLUDE = "include"
889 EXCLUDE = "exclude"
890
891
892# Type alias for pattern matching
893PatternList = list[tuple[PatternType, str]]
894
895
[docs]
896@dataclass
897class DryrunResult:
898 """
899 Holds references to JSONL files produced by a dryrun sync operation.
900
901 Each file contains one JSON object per line, matching the :py:class:`ObjectMetadata`
902 serialization format (see :py:meth:`ObjectMetadata.to_dict` / :py:meth:`ObjectMetadata.from_dict`).
903
904 The caller is responsible for cleaning up the files when they are no longer needed.
905 """
906
907 #: Path to a JSONL file listing source objects that would be added to the target.
908 files_to_add: str
909 #: Path to a JSONL file listing target objects that would be deleted.
910 files_to_delete: str
911
912
[docs]
913@dataclass
914class SyncResult:
915 """
916 A data class that represents the summary of a sync operation.
917 """
918
919 #: The total number of work units tracked for progress (including files from both source and target after filtering). Each work unit represents an ADD or DELETE operation.
920 total_work_units: int = 0
921 #: The total number of files processed to the target.
922 total_files_added: int = 0
923 #: The total number of files deleted from the target.
924 total_files_deleted: int = 0
925 #: The total number of bytes transferred to the target.
926 total_bytes_added: int = 0
927 #: The total number of bytes deleted from the target.
928 total_bytes_deleted: int = 0
929 #: The total time taken to process the sync operation.
930 total_time_seconds: float = 0.0
931 #: Dryrun details with paths to JSONL files. ``None`` for normal (non-dryrun) sync operations.
932 dryrun: DryrunResult | None = None
933
934 def __str__(self) -> str:
935 header = "Sync dryrun statistics:" if self.dryrun else "Sync statistics:"
936 lines = (
937 f"{header}\n"
938 f" Work units: {self.total_work_units}\n"
939 f" Files added: {self.total_files_added}\n"
940 f" Files deleted: {self.total_files_deleted}\n"
941 f" Bytes added: {self.total_bytes_added}\n"
942 f" Bytes deleted: {self.total_bytes_deleted}\n"
943 f" Time elapsed: {self.total_time_seconds:.2f}s"
944 )
945 if self.dryrun:
946 lines += f"\n Files to add: {self.dryrun.files_to_add}\n Files to delete: {self.dryrun.files_to_delete}"
947 return lines
948
949
[docs]
950class SyncError(RuntimeError):
951 """
952 Exception raised when errors occur during a sync operation.
953
954 This exception includes the partial SyncResult showing what was accomplished
955 before the error occurred, allowing users to understand the state of the sync.
956
957 :param message: The error message describing what went wrong.
958 :param sync_result: The partial SyncResult with statistics from the failed sync operation.
959 """
960
961 def __init__(self, message: str, sync_result: SyncResult):
962 super().__init__(message)
963 self.sync_result = sync_result
964
965 def __str__(self) -> str:
966 sync_stats = str(self.sync_result).replace("Sync statistics:", "Partial sync statistics:")
967 return f"{super().__str__()}\n\n{sync_stats}"