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