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 io
17import os
18from collections.abc import Callable, Iterator
19from datetime import UTC, datetime
20from typing import IO, Any, TypeVar
21
22from aistore.sdk import Client, RetryConfig
23from aistore.sdk.authn import AuthNClient
24from aistore.sdk.errors import AISError
25from aistore.sdk.obj.content_iterator import ParallelBuffer
26from dateutil.parser import parse as dateutil_parser
27from requests.exceptions import HTTPError
28from urllib3.util import Retry
29
30from ..constants import DEFAULT_READ_TIMEOUT
31from ..telemetry import Telemetry
32from ..types import (
33 AWARE_DATETIME_MIN,
34 Credentials,
35 CredentialsProvider,
36 ObjectMetadata,
37 Range,
38 SuffixRange,
39 SymlinkHandling,
40)
41from ..utils import safe_makedirs, split_path, validate_attributes
42from .base import BaseStorageProvider, _translate_iterator_errors
43
44_T = TypeVar("_T")
45
46PROVIDER = "ais"
47DEFAULT_PAGE_SIZE = 1000
48
49
[docs]
50class StaticAISCredentialProvider(CredentialsProvider):
51 """
52 A concrete implementation of the :py:class:`multistorageclient.types.CredentialsProvider` that provides static AIStore credentials.
53 """
54
55 _username: str | None
56 _password: str | None
57 _authn_endpoint: str | None
58 _token: str | None
59 _skip_verify: bool
60 _ca_cert: str | None
61
62 def __init__(
63 self,
64 username: str | None = None,
65 password: str | None = None,
66 authn_endpoint: str | None = None,
67 token: str | None = None,
68 skip_verify: bool = True,
69 ca_cert: str | None = None,
70 ):
71 """
72 Initializes the :py:class:`StaticAISCredentialProvider` with the given credentials.
73
74 :param username: The username for the AIStore authentication.
75 :param password: The password for the AIStore authentication.
76 :param authn_endpoint: The AIStore authentication endpoint.
77 :param token: The AIStore authentication token. This is used for authentication if username,
78 password and authn_endpoint are not provided.
79 :param skip_verify: If true, skip SSL certificate verification.
80 :param ca_cert: Path to a CA certificate file for SSL verification.
81 """
82 self._username = username
83 self._password = password
84 self._authn_endpoint = authn_endpoint
85 self._token = token
86 self._skip_verify = skip_verify
87 self._ca_cert = ca_cert
88
[docs]
89 def get_credentials(self) -> Credentials:
90 if self._username and self._password and self._authn_endpoint:
91 authn_client = AuthNClient(self._authn_endpoint, self._skip_verify, self._ca_cert)
92 self._token = authn_client.login(self._username, self._password)
93 return Credentials(token=self._token, access_key="", secret_key="", expiration=None)
94
[docs]
95 def refresh_credentials(self) -> None:
96 pass
97
98
[docs]
99class AIStoreStorageProvider(BaseStorageProvider):
100 """
101 A concrete implementation of the :py:class:`multistorageclient.types.StorageProvider` for interacting with NVIDIA AIStore.
102 """
103
104 def __init__(
105 self,
106 endpoint: str = os.getenv("AIS_ENDPOINT", ""),
107 provider: str = PROVIDER,
108 skip_verify: bool = True,
109 ca_cert: str | None = None,
110 timeout: float | tuple[float, float] | None = None,
111 retry: dict[str, Any] | None = None,
112 base_path: str = "",
113 credentials_provider: CredentialsProvider | None = None,
114 config_dict: dict[str, Any] | None = None,
115 telemetry_provider: Callable[[], Telemetry] | None = None,
116 **kwargs: Any,
117 ) -> None:
118 """
119 AIStore client for managing buckets, objects, and ETL jobs.
120
121 :param endpoint: The AIStore endpoint.
122 :param skip_verify: Whether to skip SSL certificate verification.
123 :param ca_cert: Path to a CA certificate file for SSL verification.
124 :param timeout: Request timeout in seconds; a single float
125 for both connect/read timeouts (e.g., ``5.0``), a tuple for separate connect/read
126 timeouts (e.g., ``(3.0, 10.0)``), or ``None`` (default) to use ``DEFAULT_READ_TIMEOUT`` (60 seconds).
127 :param retry: ``urllib3.util.Retry`` parameters. Applied as the ``http_retry`` of the client's
128 ``aistore.sdk.RetryConfig`` (the network-level retry defaults are preserved).
129 :param base_path: The root prefix path within the bucket where all operations will be scoped.
130 :param credentials_provider: The provider to retrieve AIStore credentials.
131 :param config_dict: Resolved MSC config.
132 :param telemetry_provider: A function that provides a telemetry instance.
133 """
134 super().__init__(
135 base_path=base_path,
136 provider_name=PROVIDER,
137 config_dict=config_dict,
138 telemetry_provider=telemetry_provider,
139 )
140
141 # https://aistore.nvidia.com/docs/python-sdk#client.Client
142 retry_config = None
143 if retry is not None:
144 retry_config = RetryConfig.default()
145 retry_config.http_retry = Retry(**retry)
146 token = None
147 if timeout is None:
148 timeout = float(DEFAULT_READ_TIMEOUT)
149 if credentials_provider:
150 token = credentials_provider.get_credentials().token
151 self.client = Client(
152 endpoint=endpoint,
153 retry_config=retry_config,
154 skip_verify=skip_verify,
155 ca_cert=ca_cert,
156 timeout=timeout,
157 token=token,
158 )
159 else:
160 self.client = Client(
161 endpoint=endpoint, retry_config=retry_config, timeout=timeout, skip_verify=skip_verify, ca_cert=ca_cert
162 )
163 self.provider = provider
164
165 def _translate_errors(
166 self,
167 func: Callable[[], _T],
168 operation: str,
169 bucket: str,
170 key: str,
171 ) -> _T:
172 """
173 Translates errors like timeouts and client errors.
174
175 :param func: The function that performs the actual object storage operation.
176 :param operation: The type of operation being performed (e.g., ``PUT``, ``GET``, ``DELETE``).
177 :param bucket: The name of the object storage bucket involved in the operation.
178 :param key: The key of the object within the object storage bucket.
179
180 :return: The result of the object storage operation, typically the return value of the `func` callable.
181 """
182
183 try:
184 return func()
185 except AISError as error:
186 status_code = error.status_code
187 if status_code == 404:
188 raise FileNotFoundError(f"Object {bucket}/{key} does not exist.") # pylint: disable=raise-missing-from
189 error_info = f"status_code: {status_code}, message: {error.message}"
190 raise RuntimeError(f"Failed to {operation} object(s) at {bucket}/{key}. {error_info}") from error
191 except HTTPError as error:
192 if error.response is not None and error.response.status_code == 404:
193 raise FileNotFoundError(f"Object {bucket}/{key} does not exist.") # pylint: disable=raise-missing-from
194 else:
195 raise RuntimeError(
196 f"Failed to {operation} object(s) at {bucket}/{key}, error type: {type(error).__name__}"
197 ) from error
198 except Exception as error:
199 raise RuntimeError(
200 f"Failed to {operation} object(s) at {bucket}/{key}, error type: {type(error).__name__}, error: {error}"
201 ) from error
202
203 def _put_object(
204 self,
205 path: str,
206 body: bytes,
207 if_match: str | None = None,
208 if_none_match: str | None = None,
209 attributes: dict[str, str] | None = None,
210 ) -> int:
211 # ais does not support if_match and if_none_match
212 bucket, key = split_path(path)
213
214 def _invoke_api() -> int:
215 obj = self.client.bucket(bucket, self.provider).object(obj_name=key)
216 validated_attributes = validate_attributes(attributes)
217 writer = obj.get_writer()
218 writer.put_content(body)
219 if validated_attributes:
220 writer.set_custom_props(custom_metadata=validated_attributes, replace_existing=True)
221
222 return len(body)
223
224 return self._translate_errors(_invoke_api, operation="PUT", bucket=bucket, key=key)
225
226 def _supports_suffix_range(self) -> bool:
227 return True
228
229 def _get_object(self, path: str, byte_range: Range | SuffixRange | None = None) -> bytes:
230 bucket, key = split_path(path)
231 if isinstance(byte_range, SuffixRange):
232 bytes_range = f"bytes=-{byte_range.size}"
233 elif byte_range:
234 bytes_range = f"bytes={byte_range.offset}-{byte_range.offset + byte_range.size - 1}"
235 else:
236 bytes_range = None
237
238 def _invoke_api() -> bytes:
239 obj = self.client.bucket(bucket, self.provider).object(obj_name=key)
240 if byte_range:
241 reader = obj.get_reader(byte_range=bytes_range) # pyright: ignore [reportArgumentType]
242 else:
243 reader = obj.get_reader()
244 content = reader.read_all()
245 if isinstance(content, bytes):
246 return content
247 elif isinstance(content, ParallelBuffer):
248 try:
249 return content.tobytes()
250 finally:
251 content.close()
252 else:
253 raise TypeError(f"Unexpected read_all() return type: {type(content)}")
254
255 if isinstance(byte_range, SuffixRange):
256 suffix_size = byte_range.size
257
258 def _invoke_suffix_api() -> bytes | None:
259 try:
260 data = _invoke_api()
261 except (AISError, HTTPError) as error:
262 # Raised for empty objects (e.g. MSC symlinks).
263 if isinstance(error, AISError):
264 status_code = error.status_code
265 else:
266 status_code = error.response.status_code if error.response is not None else None
267 if status_code == 416:
268 return None
269 raise
270 # Servers may ignore the range and return the whole object.
271 return data[-suffix_size:]
272
273 return self._get_object_suffix(
274 path,
275 byte_range,
276 lambda: self._translate_errors(_invoke_suffix_api, operation="GET", bucket=bucket, key=key),
277 )
278
279 return self._translate_errors(_invoke_api, operation="GET", bucket=bucket, key=key)
280
281 def _copy_object(self, src_path: str, dest_path: str) -> int:
282 src_bucket, src_key = split_path(src_path)
283 dest_bucket, dest_key = split_path(dest_path)
284
285 def _invoke_api() -> int:
286 src_obj = self.client.bucket(bck_name=src_bucket, provider=self.provider).object(obj_name=src_key)
287 dest_obj = self.client.bucket(bck_name=dest_bucket, provider=self.provider).object(obj_name=dest_key)
288
289 # Get source size before copying
290 src_props = src_obj.head()
291
292 # Server-side copy (preserves custom metadata automatically)
293 src_obj.copy(to_obj=dest_obj)
294
295 return int(src_props.size)
296
297 return self._translate_errors(
298 _invoke_api, operation="COPY", bucket=f"{src_bucket}->{dest_bucket}", key=f"{src_key}->{dest_key}"
299 )
300
301 def _delete_object(self, path: str, if_match: str | None = None) -> None:
302 bucket, key = split_path(path)
303
304 def _invoke_api() -> None:
305 obj = self.client.bucket(bucket, self.provider).object(obj_name=key)
306 # AIS doesn't support if-match deletion, so we implement a fallback mechanism
307 if if_match:
308 raise NotImplementedError("AIStore does not support if-match deletion")
309 # Perform deletion
310 obj.delete()
311
312 return self._translate_errors(_invoke_api, operation="DELETE", bucket=bucket, key=key)
313
314 def _is_dir(self, path: str) -> bool:
315 # Ensure the path ends with '/' to mimic a directory
316 path = self._append_delimiter(path)
317
318 bucket, prefix = split_path(path)
319
320 def _invoke_api() -> bool:
321 # List objects with the given prefix (limit to 1 for efficiency)
322 objects = self.client.bucket(bck_name=bucket, provider=self.provider).list_objects_iter(
323 prefix=prefix, page_size=1
324 )
325 # Check if there are any objects with this prefix
326 return any(True for _ in objects)
327
328 return self._translate_errors(_invoke_api, operation="LIST", bucket=bucket, key=prefix)
329
330 def _make_symlink(self, path: str, target: str) -> None:
331 bucket, key = split_path(path)
332 target_bucket, target_key = split_path(target)
333 if bucket != target_bucket:
334 raise ValueError(f"Cannot create cross-bucket symlink: '{bucket}' -> '{target_bucket}'.")
335 relative_target = ObjectMetadata.encode_symlink_target(key, target_key)
336
337 def _invoke_api() -> None:
338 obj = self.client.bucket(bucket, self.provider).object(obj_name=key)
339 writer = obj.get_writer()
340 writer.put_content(b"")
341 writer.set_custom_props(custom_metadata={"msc-symlink-target": relative_target}, replace_existing=True)
342
343 self._translate_errors(_invoke_api, operation="PUT", bucket=bucket, key=key)
344
345 def _get_object_metadata(self, path: str, strict: bool = True) -> ObjectMetadata:
346 bucket, key = split_path(path)
347 if path.endswith("/") or (bucket and not key):
348 # If path ends with "/" or empty key name is provided, then assume it's a "directory",
349 # which metadata is not guaranteed to exist for cases such as
350 # "virtual prefix" that was never explicitly created.
351 if self._is_dir(path):
352 return ObjectMetadata(
353 key=path,
354 type="directory",
355 content_length=0,
356 last_modified=AWARE_DATETIME_MIN,
357 )
358 else:
359 raise FileNotFoundError(f"Directory {path} does not exist.")
360 else:
361
362 def _invoke_api() -> ObjectMetadata:
363 obj = self.client.bucket(bck_name=bucket, provider=self.provider).object(obj_name=key)
364 try:
365 props = obj.head(props="size,atime,checksum,custom")
366
367 # The access time is not always present in the response.
368 if props.access_time:
369 last_modified = datetime.fromtimestamp(int(props.access_time) / 1e9).astimezone(UTC)
370 else:
371 last_modified = AWARE_DATETIME_MIN
372
373 user_metadata = props.custom_metadata
374 symlink_target = user_metadata.get("msc-symlink-target") if user_metadata else None
375 return ObjectMetadata(
376 key=key,
377 content_length=props.size,
378 last_modified=last_modified,
379 etag=props.checksum_value,
380 metadata=user_metadata,
381 symlink_target=symlink_target,
382 )
383 except (AISError, HTTPError) as e:
384 # Check if this might be a virtual directory (prefix with objects under it)
385 status_code = None
386 if isinstance(e, AISError):
387 status_code = e.status_code
388 elif isinstance(e, HTTPError) and e.response is not None:
389 status_code = e.response.status_code
390
391 if status_code == 404 and self._is_dir(path):
392 return ObjectMetadata(
393 key=path + "/",
394 type="directory",
395 content_length=0,
396 last_modified=AWARE_DATETIME_MIN,
397 )
398 # Re-raise to be handled by _translate_errors
399 raise
400
401 return self._translate_errors(_invoke_api, operation="HEAD", bucket=bucket, key=key)
402
403 def _list_objects(
404 self,
405 path: str,
406 start_after: str | None = None,
407 end_at: str | None = None,
408 include_directories: bool = False,
409 symlink_handling: SymlinkHandling = SymlinkHandling.FOLLOW,
410 ) -> Iterator[ObjectMetadata]:
411 bucket, prefix = split_path(path)
412
413 # Get the prefix of the start_after and end_at paths relative to the bucket.
414 if start_after:
415 _, start_after = split_path(start_after)
416 if end_at:
417 _, end_at = split_path(end_at)
418
419 def _invoke_api() -> Iterator[ObjectMetadata]:
420 # AIS has no start key option like other object stores.
421 all_objects = self.client.bucket(bck_name=bucket, provider=self.provider).list_objects_iter(
422 prefix=prefix, props="name,size,atime,checksum,cone", page_size=DEFAULT_PAGE_SIZE
423 )
424
425 # Assume AIS guarantees lexicographical order.
426 for bucket_entry in all_objects:
427 obj = bucket_entry.object
428 key = obj.name
429 props = bucket_entry.generate_object_props()
430
431 # The access time is not always present in the response.
432 if props.access_time:
433 last_modified = dateutil_parser(props.access_time).astimezone(UTC)
434 else:
435 last_modified = AWARE_DATETIME_MIN
436
437 if (start_after is None or start_after < key) and (end_at is None or key <= end_at):
438 yield ObjectMetadata(
439 key=key, content_length=int(props.size), last_modified=last_modified, etag=props.checksum_value
440 )
441 elif end_at is not None and end_at < key:
442 return
443
444 return _translate_iterator_errors(
445 _invoke_api(), lambda f: self._translate_errors(f, operation="LIST", bucket=bucket, key=prefix)
446 )
447
448 def _upload_file(self, remote_path: str, f: str | IO, attributes: dict[str, str] | None = None) -> int:
449 file_size: int = 0
450
451 if isinstance(f, str):
452 with open(f, "rb") as fp:
453 body = fp.read()
454 file_size = len(body)
455 self._put_object(remote_path, body, attributes=attributes)
456 else:
457 if isinstance(f, io.StringIO):
458 body = f.read().encode("utf-8")
459 file_size = len(body)
460 self._put_object(remote_path, body, attributes=attributes)
461 else:
462 body = f.read()
463 file_size = len(body)
464 self._put_object(remote_path, body, attributes=attributes)
465
466 return file_size
467
468 def _download_file(self, remote_path: str, f: str | IO, metadata: ObjectMetadata | None = None) -> int:
469 if metadata is None:
470 metadata = self._get_object_metadata(remote_path)
471
472 if isinstance(f, str):
473 if os.path.dirname(f):
474 safe_makedirs(os.path.dirname(f))
475 with open(f, "wb") as fp:
476 fp.write(self._get_object(remote_path))
477 else:
478 if isinstance(f, io.StringIO):
479 f.write(self._get_object(remote_path).decode("utf-8"))
480 else:
481 f.write(self._get_object(remote_path))
482
483 return metadata.content_length