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