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, RetryConfig
23from aistore.sdk.authn import AuthNClient
24from aistore.sdk.errors import AISError
25from aistore.sdk.obj.content_iterator import ParallelBuffer
26from aistore.sdk.obj.object_props import ObjectProps
27from dateutil.parser import parse as dateutil_parser
28from requests.exceptions import HTTPError
29from urllib3.util import Retry
30
31from ..constants import DEFAULT_READ_TIMEOUT
32from ..telemetry import Telemetry
33from ..types import (
34 AWARE_DATETIME_MIN,
35 Credentials,
36 CredentialsProvider,
37 ObjectMetadata,
38 Range,
39 SymlinkHandling,
40)
41from ..utils import safe_makedirs, split_path, validate_attributes
42from .base import BaseStorageProvider
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: Optional[str]
56 _password: Optional[str]
57 _authn_endpoint: Optional[str]
58 _token: Optional[str]
59 _skip_verify: bool
60 _ca_cert: Optional[str]
61
62 def __init__(
63 self,
64 username: Optional[str] = None,
65 password: Optional[str] = None,
66 authn_endpoint: Optional[str] = None,
67 token: Optional[str] = None,
68 skip_verify: bool = True,
69 ca_cert: Optional[str] = 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: Optional[str] = None,
110 timeout: Optional[Union[float, tuple[float, float]]] = None,
111 retry: Optional[dict[str, Any]] = None,
112 base_path: str = "",
113 credentials_provider: Optional[CredentialsProvider] = None,
114 config_dict: Optional[dict[str, Any]] = None,
115 telemetry_provider: Optional[Callable[[], Telemetry]] = 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`` to disable timeout.
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 status_code = error.response.status_code
193 if status_code == 404:
194 raise FileNotFoundError(f"Object {bucket}/{key} does not exist.") # pylint: disable=raise-missing-from
195 else:
196 raise RuntimeError(
197 f"Failed to {operation} object(s) at {bucket}/{key}, error type: {type(error).__name__}"
198 ) from error
199 except Exception as error:
200 raise RuntimeError(
201 f"Failed to {operation} object(s) at {bucket}/{key}, error type: {type(error).__name__}, error: {error}"
202 ) from error
203
204 def _put_object(
205 self,
206 path: str,
207 body: bytes,
208 if_match: Optional[str] = None,
209 if_none_match: Optional[str] = None,
210 attributes: Optional[dict[str, str]] = None,
211 ) -> int:
212 # ais does not support if_match and if_none_match
213 bucket, key = split_path(path)
214
215 def _invoke_api() -> int:
216 obj = self.client.bucket(bucket, self.provider).object(obj_name=key)
217 validated_attributes = validate_attributes(attributes)
218 writer = obj.get_writer()
219 writer.put_content(body)
220 if validated_attributes:
221 writer.set_custom_props(custom_metadata=validated_attributes, replace_existing=True)
222
223 return len(body)
224
225 return self._translate_errors(_invoke_api, operation="PUT", bucket=bucket, key=key)
226
227 def _get_object(self, path: str, byte_range: Optional[Range] = None) -> bytes:
228 bucket, key = split_path(path)
229 if byte_range:
230 bytes_range = f"bytes={byte_range.offset}-{byte_range.offset + byte_range.size - 1}"
231 else:
232 bytes_range = None
233
234 def _invoke_api() -> bytes:
235 obj = self.client.bucket(bucket, self.provider).object(obj_name=key)
236 if byte_range:
237 reader = obj.get_reader(byte_range=bytes_range) # pyright: ignore [reportArgumentType]
238 else:
239 reader = obj.get_reader()
240 content = reader.read_all()
241 if isinstance(content, bytes):
242 return content
243 elif isinstance(content, ParallelBuffer):
244 try:
245 return content.tobytes()
246 finally:
247 content.close()
248 else:
249 raise TypeError(f"Unexpected read_all() return type: {type(content)}")
250
251 return self._translate_errors(_invoke_api, operation="GET", bucket=bucket, key=key)
252
253 def _copy_object(self, src_path: str, dest_path: str) -> int:
254 src_bucket, src_key = split_path(src_path)
255 dest_bucket, dest_key = split_path(dest_path)
256
257 def _invoke_api() -> int:
258 src_obj = self.client.bucket(bck_name=src_bucket, provider=self.provider).object(obj_name=src_key)
259 dest_obj = self.client.bucket(bck_name=dest_bucket, provider=self.provider).object(obj_name=dest_key)
260
261 # Get source size before copying
262 src_headers = src_obj.head()
263 src_props = ObjectProps(src_headers)
264
265 # Server-side copy (preserves custom metadata automatically)
266 src_obj.copy(to_obj=dest_obj) # type: ignore[attr-defined]
267
268 return int(src_props.size)
269
270 return self._translate_errors(
271 _invoke_api, operation="COPY", bucket=f"{src_bucket}->{dest_bucket}", key=f"{src_key}->{dest_key}"
272 )
273
274 def _delete_object(self, path: str, if_match: Optional[str] = None) -> None:
275 bucket, key = split_path(path)
276
277 def _invoke_api() -> None:
278 obj = self.client.bucket(bucket, self.provider).object(obj_name=key)
279 # AIS doesn't support if-match deletion, so we implement a fallback mechanism
280 if if_match:
281 raise NotImplementedError("AIStore does not support if-match deletion")
282 # Perform deletion
283 obj.delete()
284
285 return self._translate_errors(_invoke_api, operation="DELETE", bucket=bucket, key=key)
286
287 def _is_dir(self, path: str) -> bool:
288 # Ensure the path ends with '/' to mimic a directory
289 path = self._append_delimiter(path)
290
291 bucket, prefix = split_path(path)
292
293 def _invoke_api() -> bool:
294 # List objects with the given prefix (limit to 1 for efficiency)
295 objects = self.client.bucket(bck_name=bucket, provider=self.provider).list_objects_iter(
296 prefix=prefix, page_size=1
297 )
298 # Check if there are any objects with this prefix
299 return any(True for _ in objects)
300
301 return self._translate_errors(_invoke_api, operation="LIST", bucket=bucket, key=prefix)
302
303 def _make_symlink(self, path: str, target: str) -> None:
304 bucket, key = split_path(path)
305 target_bucket, target_key = split_path(target)
306 if bucket != target_bucket:
307 raise ValueError(f"Cannot create cross-bucket symlink: '{bucket}' -> '{target_bucket}'.")
308 relative_target = ObjectMetadata.encode_symlink_target(key, target_key)
309
310 def _invoke_api() -> None:
311 obj = self.client.bucket(bucket, self.provider).object(obj_name=key)
312 writer = obj.get_writer()
313 writer.put_content(b"")
314 writer.set_custom_props(custom_metadata={"msc-symlink-target": relative_target}, replace_existing=True)
315
316 self._translate_errors(_invoke_api, operation="PUT", bucket=bucket, key=key)
317
318 def _get_object_metadata(self, path: str, strict: bool = True) -> ObjectMetadata:
319 bucket, key = split_path(path)
320 if path.endswith("/") or (bucket and not key):
321 # If path ends with "/" or empty key name is provided, then assume it's a "directory",
322 # which metadata is not guaranteed to exist for cases such as
323 # "virtual prefix" that was never explicitly created.
324 if self._is_dir(path):
325 return ObjectMetadata(
326 key=path,
327 type="directory",
328 content_length=0,
329 last_modified=AWARE_DATETIME_MIN,
330 )
331 else:
332 raise FileNotFoundError(f"Directory {path} does not exist.")
333 else:
334
335 def _invoke_api() -> ObjectMetadata:
336 obj = self.client.bucket(bck_name=bucket, provider=self.provider).object(obj_name=key)
337 try:
338 headers = obj.head()
339 props = ObjectProps(headers)
340
341 # The access time is not always present in the response.
342 if props.access_time:
343 last_modified = datetime.fromtimestamp(int(props.access_time) / 1e9).astimezone(timezone.utc)
344 else:
345 last_modified = AWARE_DATETIME_MIN
346
347 user_metadata = props.custom_metadata
348 symlink_target = user_metadata.get("msc-symlink-target") if user_metadata else None
349 return ObjectMetadata(
350 key=key,
351 content_length=int(props.size), # pyright: ignore [reportArgumentType]
352 last_modified=last_modified,
353 etag=props.checksum_value,
354 metadata=user_metadata,
355 symlink_target=symlink_target,
356 )
357 except (AISError, HTTPError) as e:
358 # Check if this might be a virtual directory (prefix with objects under it)
359 status_code = None
360 if isinstance(e, AISError):
361 status_code = e.status_code
362 elif isinstance(e, HTTPError):
363 status_code = e.response.status_code
364
365 if status_code == 404:
366 if self._is_dir(path):
367 return ObjectMetadata(
368 key=path + "/",
369 type="directory",
370 content_length=0,
371 last_modified=AWARE_DATETIME_MIN,
372 )
373 # Re-raise to be handled by _translate_errors
374 raise
375
376 return self._translate_errors(_invoke_api, operation="HEAD", bucket=bucket, key=key)
377
378 def _list_objects(
379 self,
380 path: str,
381 start_after: Optional[str] = None,
382 end_at: Optional[str] = None,
383 include_directories: bool = False,
384 symlink_handling: SymlinkHandling = SymlinkHandling.FOLLOW,
385 ) -> Iterator[ObjectMetadata]:
386 bucket, prefix = split_path(path)
387
388 # Get the prefix of the start_after and end_at paths relative to the bucket.
389 if start_after:
390 _, start_after = split_path(start_after)
391 if end_at:
392 _, end_at = split_path(end_at)
393
394 def _invoke_api() -> Iterator[ObjectMetadata]:
395 # AIS has no start key option like other object stores.
396 all_objects = self.client.bucket(bck_name=bucket, provider=self.provider).list_objects_iter(
397 prefix=prefix, props="name,size,atime,checksum,cone", page_size=DEFAULT_PAGE_SIZE
398 )
399
400 # Assume AIS guarantees lexicographical order.
401 for bucket_entry in all_objects:
402 obj = bucket_entry.object
403 key = obj.name
404 props = bucket_entry.generate_object_props()
405
406 # The access time is not always present in the response.
407 if props.access_time:
408 last_modified = dateutil_parser(props.access_time).astimezone(timezone.utc)
409 else:
410 last_modified = AWARE_DATETIME_MIN
411
412 if (start_after is None or start_after < key) and (end_at is None or key <= end_at):
413 yield ObjectMetadata(
414 key=key, content_length=int(props.size), last_modified=last_modified, etag=props.checksum_value
415 )
416 elif end_at is not None and end_at < key:
417 return
418
419 return self._translate_errors(_invoke_api, operation="LIST", bucket=bucket, key=prefix)
420
421 def _upload_file(self, remote_path: str, f: Union[str, IO], attributes: Optional[dict[str, str]] = None) -> int:
422 file_size: int = 0
423
424 if isinstance(f, str):
425 with open(f, "rb") as fp:
426 body = fp.read()
427 file_size = len(body)
428 self._put_object(remote_path, body, attributes=attributes)
429 else:
430 if isinstance(f, io.StringIO):
431 body = f.read().encode("utf-8")
432 file_size = len(body)
433 self._put_object(remote_path, body, attributes=attributes)
434 else:
435 body = f.read()
436 file_size = len(body)
437 self._put_object(remote_path, body, attributes=attributes)
438
439 return file_size
440
441 def _download_file(self, remote_path: str, f: Union[str, IO], metadata: Optional[ObjectMetadata] = None) -> int:
442 if metadata is None:
443 metadata = self._get_object_metadata(remote_path)
444
445 if isinstance(f, str):
446 if os.path.dirname(f):
447 safe_makedirs(os.path.dirname(f))
448 with open(f, "wb") as fp:
449 fp.write(self._get_object(remote_path))
450 else:
451 if isinstance(f, io.StringIO):
452 f.write(self._get_object(remote_path).decode("utf-8"))
453 else:
454 f.write(self._get_object(remote_path))
455
456 return metadata.content_length