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 logging
17from collections.abc import Iterator, Sequence
18from typing import IO, Any
19
20from ..config import StorageClientConfig
21from ..constants import MEMORY_LOAD_LIMIT
22from ..file import ObjectFile, PosixFile
23from ..types import (
24 MSC_PROTOCOL,
25 BatchTransferCallback,
26 ExecutionMode,
27 MetadataProvider,
28 ObjectMetadata,
29 PatternList,
30 Range,
31 SignerType,
32 SourceVersionCheckMode,
33 SuffixRange,
34 SymlinkHandling,
35 SyncResult,
36)
37from ..utils import PatternMatcher, _invoke_batch_transfer, join_paths
38from .single import SingleStorageClient
39from .types import AbstractStorageClient
40
41logger = logging.getLogger(__name__)
42
43
[docs]
44class CompositeStorageClient(AbstractStorageClient):
45 """
46 READ-ONLY storage client for multi-backend configurations.
47
48 Routes read operations to child SingleStorageClient instances based on
49 metadata provider's routing information (ResolvedPath.profile).
50
51 Write operations raise NotImplementedError with clear error messages.
52 """
53
54 _metadata_provider: MetadataProvider
55
56 def __init__(self, config: StorageClientConfig):
57 """
58 Initialize a composite storage client.
59
60 :param config: Storage client configuration with storage_provider_profiles set
61 :raises ValueError: If config doesn't have storage_provider_profiles
62 """
63 if not config.storage_provider_profiles:
64 raise ValueError(
65 "CompositeStorageClient requires storage_provider_profiles. "
66 "Use SingleStorageClient for single-backend configurations."
67 )
68
69 if config.metadata_provider is None:
70 raise ValueError("CompositeStorageClient requires a metadata_provider for routing decisions.")
71
72 self._config = config
73 self._profile = config.profile
74 self._metadata_provider = config.metadata_provider
75 self._metadata_provider_lock = None
76 self._storage_provider = config.storage_provider
77 self._credentials_provider = config.credentials_provider
78 self._retry_config = config.retry_config
79 self._cache_manager = config.cache_manager
80 self._replica_manager = None
81
82 self._child_clients: dict[str, SingleStorageClient] = {}
83 self._child_profile_names = config.storage_provider_profiles
84
85 if config.child_configs:
86 # ProviderBundleV2 path: child configs are pre-built
87 for child_name, child_config in config.child_configs.items():
88 self._child_clients[child_name] = SingleStorageClient(child_config)
89 else:
90 # Config-dict path: child profiles have their own credentials defined
91 if not config._config_dict:
92 raise ValueError("CompositeStorageClient requires _config_dict to build child clients")
93
94 for child_profile in self._child_profile_names:
95 child_config = StorageClientConfig.from_dict(
96 config_dict=config._config_dict,
97 profile=child_profile,
98 )
99 self._child_clients[child_profile] = SingleStorageClient(child_config)
100
101 @property
102 def profile(self) -> str:
103 """
104 :return: The profile name of the storage client.
105 """
106 return self._profile
107
108 @property
109 def replicas(self) -> list[AbstractStorageClient]:
110 """
111 :return: List of replica storage clients (empty list for CompositeStorageClient).
112 """
113 return []
114
[docs]
115 def is_default_profile(self) -> bool:
116 return self._config.profile == "__filesystem__"
117
118 def _is_rust_client_enabled(self) -> bool:
119 """
120 Check if Rust client is enabled for all child storage clients.
121
122 When all child backends are Rust-enabled, MSC can use single-process
123 multi-threaded mode instead of multi-process mode, as Rust handles
124 parallelism internally via async I/O without Python GIL contention.
125
126 :return: True if all child clients are Rust-enabled, False otherwise.
127 """
128 return all(client._is_rust_client_enabled() for client in self._child_clients.values())
129
130 def _is_posix_file_storage_provider(self) -> bool:
131 """
132 Check if using POSIX file storage provider.
133
134 :return: False - composite client doesn't have a single storage provider.
135 """
136 return False
137
[docs]
138 def get_posix_path(self, path: str) -> str | None:
139 """
140 Get the POSIX filesystem path for a given logical path.
141
142 :param path: The logical path to resolve.
143 :return: None - composite client doesn't support POSIX path resolution.
144 """
145 return None
146
147 def _get_child_client(self, profile: str | None) -> SingleStorageClient:
148 """
149 Get the child client for the specified profile.
150
151 :param profile: Profile name from ResolvedPath
152 :return: Child SingleStorageClient instance
153 :raises ValueError: If profile is None or not found
154 """
155 if profile is None:
156 raise ValueError(
157 "CompositeStorageClient requires profile from ResolvedPath for routing. "
158 "Metadata provider must return ResolvedPath with profile set."
159 )
160
161 if profile not in self._child_clients:
162 raise ValueError(
163 f"Profile '{profile}' not found in composite client. "
164 f"Available profiles: {list(self._child_clients.keys())}"
165 )
166
167 return self._child_clients[profile]
168
[docs]
169 def read(
170 self,
171 path: str,
172 byte_range: Range | SuffixRange | None = None,
173 check_source_version: SourceVersionCheckMode = SourceVersionCheckMode.INHERIT,
174 ) -> bytes:
175 resolved = self._metadata_provider.realpath(path)
176 if not resolved.exists:
177 raise FileNotFoundError(f"Path '{path}' not found")
178
179 child = self._get_child_client(resolved.profile)
180 return child.read(resolved.physical_path, byte_range, check_source_version)
181
[docs]
182 def open(
183 self,
184 path: str,
185 mode: str = "rb",
186 buffering: int = -1,
187 encoding: str | None = None,
188 disable_read_cache: bool = False,
189 memory_load_limit: int = MEMORY_LOAD_LIMIT,
190 atomic: bool = True,
191 check_source_version: SourceVersionCheckMode = SourceVersionCheckMode.INHERIT,
192 attributes: dict[str, Any] | None = None,
193 prefetch_file: bool | None = None,
194 ) -> PosixFile | ObjectFile:
195 """
196 Open a file for reading or writing.
197
198 :param path: The logical path of the object to open.
199 :param mode: The file mode. Supported modes: "r", "rb", "w", "wb", "a", "ab".
200 :param buffering: The buffering mode. Only applies to PosixFile.
201 :param encoding: The encoding to use for text files.
202 :param disable_read_cache: When set to ``True``, disables caching for file content.
203 This parameter is only applicable to ObjectFile when the mode is "r" or "rb".
204 :param memory_load_limit: Size limit in bytes for loading files into memory. Defaults to 512MB.
205 This parameter is only applicable to ObjectFile when the mode is "r" or "rb". Defaults to 512MB.
206 :param atomic: When set to ``True``, file will be written atomically (rename upon close).
207 This parameter is only applicable to PosixFile in write mode.
208 :param check_source_version: Whether to check the source version of cached objects.
209 :param attributes: Attributes to add to the file.
210 This parameter is only applicable when the mode is "w" or "wb" or "a" or "ab". Defaults to None.
211 :param prefetch_file: Whether to prefetch the file content.
212 This parameter is only applicable to ObjectFile when the mode is "r" or "rb".
213 If None, inherits from cache configuration.
214 :return: A file-like object (PosixFile or ObjectFile) for the specified path.
215 :raises FileNotFoundError: If the file does not exist (read mode).
216 :raises NotImplementedError: If the operation is not supported (e.g., write on CompositeStorageClient).
217 """
218 if mode not in ["r", "rb"]:
219 raise NotImplementedError(
220 f"CompositeStorageClient only supports read mode (got '{mode}'). "
221 "Write operations are not implemented for multi-location datasets."
222 )
223
224 resolved = self._metadata_provider.realpath(path)
225 if not resolved.exists:
226 raise FileNotFoundError(f"Path '{path}' not found")
227
228 child = self._get_child_client(resolved.profile)
229 return child.open(
230 resolved.physical_path,
231 mode,
232 buffering,
233 encoding,
234 disable_read_cache,
235 memory_load_limit,
236 atomic,
237 check_source_version,
238 attributes,
239 prefetch_file,
240 )
241
[docs]
242 def download_file(self, remote_path: str, local_path: str | IO, metadata: ObjectMetadata | None = None) -> None:
243 resolved = self._metadata_provider.realpath(remote_path)
244 if not resolved.exists:
245 raise FileNotFoundError(f"Path '{remote_path}' not found")
246
247 child = self._get_child_client(resolved.profile)
248 child.download_file(resolved.physical_path, local_path, metadata)
249
[docs]
250 def download_files(
251 self,
252 remote_paths: Sequence[str],
253 local_paths: Sequence[str],
254 metadata: Sequence[ObjectMetadata | None] | None = None,
255 max_workers: int = 16,
256 transfer_callback: BatchTransferCallback | None = None,
257 ) -> None:
258 if len(remote_paths) != len(local_paths):
259 raise ValueError("remote_paths and local_paths must have the same length")
260
261 groups: dict[str, tuple[list[str], list[str], list[ObjectMetadata | None]]] = {}
262 for i, (remote_path, local_path) in enumerate(zip(remote_paths, local_paths)):
263 resolved = self._metadata_provider.realpath(remote_path)
264 if not resolved.exists:
265 raise FileNotFoundError(f"Path '{remote_path}' not found")
266
267 profile = resolved.profile
268 if profile is None:
269 raise ValueError(
270 "CompositeStorageClient requires profile from ResolvedPath for routing. "
271 "Metadata provider must return ResolvedPath with profile set."
272 )
273 if profile not in groups:
274 groups[profile] = ([], [], [])
275 groups[profile][0].append(resolved.physical_path)
276 groups[profile][1].append(local_path)
277 groups[profile][2].append(metadata[i] if metadata is not None else None)
278
279 for profile, (physical_paths, group_local_paths, group_metadata) in groups.items():
280 child = self._get_child_client(profile)
281 has_any_metadata = any(m is not None for m in group_metadata)
282 _invoke_batch_transfer(
283 child.download_files,
284 physical_paths,
285 group_local_paths,
286 group_metadata if has_any_metadata else None,
287 max_workers,
288 transfer_callback=transfer_callback,
289 )
290
[docs]
291 def glob(
292 self,
293 pattern: str,
294 include_url_prefix: bool = False,
295 attribute_filter_expression: str | None = None,
296 ) -> list[str]:
297 results = self._metadata_provider.glob(
298 pattern,
299 attribute_filter_expression=attribute_filter_expression,
300 )
301
302 if include_url_prefix:
303 results = [join_paths(f"{MSC_PROTOCOL}{self._config.profile}", path) for path in results]
304
305 return results
306
[docs]
307 def list_recursive(
308 self,
309 path: str = "",
310 start_after: str | None = None,
311 end_at: str | None = None,
312 max_workers: int = 32,
313 look_ahead: int = 2,
314 include_url_prefix: bool = False,
315 patterns: PatternList | None = None,
316 symlink_handling: SymlinkHandling = SymlinkHandling.FOLLOW,
317 ) -> Iterator[ObjectMetadata]:
318 yield from self.list(
319 path=path,
320 start_after=start_after,
321 end_at=end_at,
322 include_directories=False,
323 include_url_prefix=include_url_prefix,
324 patterns=patterns,
325 symlink_handling=symlink_handling,
326 )
327
[docs]
328 def is_file(self, path: str) -> bool:
329 resolved = self._metadata_provider.realpath(path)
330 return resolved.exists
331
[docs]
332 def is_empty(self, path: str) -> bool:
333 objects = self._metadata_provider.list_objects(path)
334
335 try:
336 return next(objects) is None
337 except StopIteration:
338 pass
339
340 return True
341
[docs]
342 def info(self, path: str, strict: bool = True) -> ObjectMetadata:
343 return self._metadata_provider.get_object_metadata(path, include_pending=not strict)
344
[docs]
345 def write(
346 self,
347 path: str,
348 body: bytes,
349 attributes: dict[str, Any] | None = None,
350 ) -> None:
351 """Write operations not supported in read-only mode."""
352 raise NotImplementedError(
353 "CompositeStorageClient is read-only. "
354 "Write operations are not implemented for multi-location datasets. "
355 "Use a single-location dataset for write operations."
356 )
357
[docs]
358 def delete(self, path: str, recursive: bool = False) -> None:
359 """Delete operations not supported in read-only mode."""
360 raise NotImplementedError(
361 "CompositeStorageClient is read-only. Delete operations are not implemented for multi-location datasets."
362 )
363
[docs]
364 def delete_many(self, paths: Sequence[str]) -> None:
365 """Delete operations not supported in read-only mode."""
366 raise NotImplementedError(
367 "CompositeStorageClient is read-only. Delete operations are not implemented for multi-location datasets."
368 )
369
[docs]
370 def copy(self, src_path: str, dest_path: str) -> None:
371 """Copy operations not supported in read-only mode."""
372 raise NotImplementedError(
373 "CompositeStorageClient is read-only. Copy operations are not implemented for multi-location datasets."
374 )
375
[docs]
376 def make_symlink(self, path: str, target: str) -> None:
377 """Symlink operations not supported in read-only mode."""
378 raise NotImplementedError(
379 "CompositeStorageClient is read-only. Symlink operations are not implemented for multi-location datasets."
380 )
381
[docs]
382 def upload_file(
383 self,
384 remote_path: str,
385 local_path: str | IO,
386 attributes: dict[str, Any] | None = None,
387 ) -> None:
388 """Upload operations not supported in read-only mode."""
389 raise NotImplementedError(
390 "CompositeStorageClient is read-only. Upload operations are not implemented for multi-location datasets."
391 )
392
[docs]
393 def upload_files(
394 self,
395 remote_paths: Sequence[str],
396 local_paths: Sequence[str],
397 attributes: Sequence[dict[str, Any] | None] | None = None,
398 max_workers: int = 16,
399 transfer_callback: BatchTransferCallback | None = None,
400 ) -> None:
401 """Upload operations not supported in read-only mode."""
402 raise NotImplementedError(
403 "CompositeStorageClient is read-only. Upload operations are not implemented for multi-location datasets."
404 )
405
408
[docs]
409 def sync_from(
410 self,
411 source_client: AbstractStorageClient,
412 source_path: str = "",
413 target_path: str = "",
414 delete_unmatched_files: bool = False,
415 description: str = "Syncing",
416 num_worker_processes: int | None = None,
417 execution_mode: ExecutionMode = ExecutionMode.LOCAL,
418 patterns: PatternList | None = None,
419 preserve_source_attributes: bool = False,
420 source_files: list[str] | None = None,
421 ignore_hidden: bool = True,
422 commit_metadata: bool = True,
423 dryrun: bool = False,
424 dryrun_output_path: str | None = None,
425 symlink_handling: SymlinkHandling = SymlinkHandling.FOLLOW,
426 ) -> SyncResult:
427 raise NotImplementedError(
428 "CompositeStorageClient cannot be used as sync target (write operation). "
429 "Use CompositeStorageClient as source only, or use a single-location dataset as target."
430 )
431
[docs]
432 def sync_replicas(
433 self,
434 source_path: str,
435 replica_indices: Sequence[int] | None = None,
436 delete_unmatched_files: bool = False,
437 description: str = "Syncing replica",
438 num_worker_processes: int | None = None,
439 execution_mode: ExecutionMode = ExecutionMode.LOCAL,
440 patterns: PatternList | None = None,
441 ignore_hidden: bool = True,
442 symlink_handling: SymlinkHandling = SymlinkHandling.FOLLOW,
443 ) -> None:
444 """No-op for read-only client."""
445
[docs]
446 def list(
447 self,
448 path: str = "",
449 start_after: str | None = None,
450 end_at: str | None = None,
451 include_directories: bool = False,
452 include_url_prefix: bool = False,
453 attribute_filter_expression: str | None = None,
454 show_attributes: bool = False,
455 patterns: PatternList | None = None,
456 symlink_handling: SymlinkHandling = SymlinkHandling.FOLLOW,
457 ) -> Iterator[ObjectMetadata]:
458 # Apply patterns to the objects
459 pattern_matcher = PatternMatcher(patterns) if patterns else None
460
461 # Delegate to metadata provider (always present for CompositeStorageClient)
462 for obj in self._metadata_provider.list_objects(
463 path,
464 start_after=start_after,
465 end_at=end_at,
466 include_directories=include_directories,
467 attribute_filter_expression=attribute_filter_expression,
468 show_attributes=show_attributes,
469 ):
470 # Skip objects that do not match the patterns
471 if pattern_matcher and not pattern_matcher.should_include_file(obj.key):
472 continue
473
474 if include_url_prefix:
475 # Return a copy: the object is owned by the metadata provider and must not be mutated.
476 obj = obj.replace(key=join_paths(f"{MSC_PROTOCOL}{self._config.profile}", obj.key))
477
478 yield obj
479
[docs]
480 def generate_presigned_url(
481 self,
482 path: str,
483 *,
484 method: str = "GET",
485 signer_type: SignerType | None = None,
486 signer_options: dict[str, Any] | None = None,
487 ) -> str:
488 resolved = self._metadata_provider.realpath(path)
489 if not resolved.exists:
490 raise FileNotFoundError(f"Path '{path}' not found")
491
492 child = self._get_child_client(resolved.profile)
493 return child.generate_presigned_url(
494 resolved.physical_path, method=method, signer_type=signer_type, signer_options=signer_options
495 )
496
497 def __getstate__(self) -> dict[str, Any]:
498 """Support for pickling."""
499 return {"_config": self._config}
500
501 def __setstate__(self, state: dict[str, Any]) -> None:
502 """Support for unpickling - reinitialize from config."""
503 config = state["_config"]
504 self.__init__(config)