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