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 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: 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 _get_object(self, path: str, byte_range: Range | None = None) -> bytes:
227 bucket, key = split_path(path)
228 if byte_range:
229 bytes_range = f"bytes={byte_range.offset}-{byte_range.offset + byte_range.size - 1}"
230 else:
231 bytes_range = None
232
233 def _invoke_api() -> bytes:
234 obj = self.client.bucket(bucket, self.provider).object(obj_name=key)
235 if byte_range:
236 reader = obj.get_reader(byte_range=bytes_range) # pyright: ignore [reportArgumentType]
237 else:
238 reader = obj.get_reader()
239 content = reader.read_all()
240 if isinstance(content, bytes):
241 return content
242 elif isinstance(content, ParallelBuffer):
243 try:
244 return content.tobytes()
245 finally:
246 content.close()
247 else:
248 raise TypeError(f"Unexpected read_all() return type: {type(content)}")
249
250 return self._translate_errors(_invoke_api, operation="GET", bucket=bucket, key=key)
251
252 def _copy_object(self, src_path: str, dest_path: str) -> int:
253 src_bucket, src_key = split_path(src_path)
254 dest_bucket, dest_key = split_path(dest_path)
255
256 def _invoke_api() -> int:
257 src_obj = self.client.bucket(bck_name=src_bucket, provider=self.provider).object(obj_name=src_key)
258 dest_obj = self.client.bucket(bck_name=dest_bucket, provider=self.provider).object(obj_name=dest_key)
259
260 # Get source size before copying
261 src_headers = src_obj.head()
262 src_props = ObjectProps(src_headers)
263
264 # Server-side copy (preserves custom metadata automatically)
265 src_obj.copy(to_obj=dest_obj) # type: ignore[attr-defined]
266
267 return int(src_props.size)
268
269 return self._translate_errors(
270 _invoke_api, operation="COPY", bucket=f"{src_bucket}->{dest_bucket}", key=f"{src_key}->{dest_key}"
271 )
272
273 def _delete_object(self, path: str, if_match: str | None = None) -> None:
274 bucket, key = split_path(path)
275
276 def _invoke_api() -> None:
277 obj = self.client.bucket(bucket, self.provider).object(obj_name=key)
278 # AIS doesn't support if-match deletion, so we implement a fallback mechanism
279 if if_match:
280 raise NotImplementedError("AIStore does not support if-match deletion")
281 # Perform deletion
282 obj.delete()
283
284 return self._translate_errors(_invoke_api, operation="DELETE", bucket=bucket, key=key)
285
286 def _is_dir(self, path: str) -> bool:
287 # Ensure the path ends with '/' to mimic a directory
288 path = self._append_delimiter(path)
289
290 bucket, prefix = split_path(path)
291
292 def _invoke_api() -> bool:
293 # List objects with the given prefix (limit to 1 for efficiency)
294 objects = self.client.bucket(bck_name=bucket, provider=self.provider).list_objects_iter(
295 prefix=prefix, page_size=1
296 )
297 # Check if there are any objects with this prefix
298 return any(True for _ in objects)
299
300 return self._translate_errors(_invoke_api, operation="LIST", bucket=bucket, key=prefix)
301
302 def _make_symlink(self, path: str, target: str) -> None:
303 bucket, key = split_path(path)
304 target_bucket, target_key = split_path(target)
305 if bucket != target_bucket:
306 raise ValueError(f"Cannot create cross-bucket symlink: '{bucket}' -> '{target_bucket}'.")
307 relative_target = ObjectMetadata.encode_symlink_target(key, target_key)
308
309 def _invoke_api() -> None:
310 obj = self.client.bucket(bucket, self.provider).object(obj_name=key)
311 writer = obj.get_writer()
312 writer.put_content(b"")
313 writer.set_custom_props(custom_metadata={"msc-symlink-target": relative_target}, replace_existing=True)
314
315 self._translate_errors(_invoke_api, operation="PUT", bucket=bucket, key=key)
316
317 def _get_object_metadata(self, path: str, strict: bool = True) -> ObjectMetadata:
318 bucket, key = split_path(path)
319 if path.endswith("/") or (bucket and not key):
320 # If path ends with "/" or empty key name is provided, then assume it's a "directory",
321 # which metadata is not guaranteed to exist for cases such as
322 # "virtual prefix" that was never explicitly created.
323 if self._is_dir(path):
324 return ObjectMetadata(
325 key=path,
326 type="directory",
327 content_length=0,
328 last_modified=AWARE_DATETIME_MIN,
329 )
330 else:
331 raise FileNotFoundError(f"Directory {path} does not exist.")
332 else:
333
334 def _invoke_api() -> ObjectMetadata:
335 obj = self.client.bucket(bck_name=bucket, provider=self.provider).object(obj_name=key)
336 try:
337 headers = obj.head()
338 props = ObjectProps(headers)
339
340 # The access time is not always present in the response.
341 if props.access_time:
342 last_modified = datetime.fromtimestamp(int(props.access_time) / 1e9).astimezone(timezone.utc)
343 else:
344 last_modified = AWARE_DATETIME_MIN
345
346 user_metadata = props.custom_metadata
347 symlink_target = user_metadata.get("msc-symlink-target") if user_metadata else None
348 return ObjectMetadata(
349 key=key,
350 content_length=int(props.size), # pyright: ignore [reportArgumentType]
351 last_modified=last_modified,
352 etag=props.checksum_value,
353 metadata=user_metadata,
354 symlink_target=symlink_target,
355 )
356 except (AISError, HTTPError) as e:
357 # Check if this might be a virtual directory (prefix with objects under it)
358 status_code = None
359 if isinstance(e, AISError):
360 status_code = e.status_code
361 elif isinstance(e, HTTPError) and e.response is not None:
362 status_code = e.response.status_code
363
364 if status_code == 404 and self._is_dir(path):
365 return ObjectMetadata(
366 key=path + "/",
367 type="directory",
368 content_length=0,
369 last_modified=AWARE_DATETIME_MIN,
370 )
371 # Re-raise to be handled by _translate_errors
372 raise
373
374 return self._translate_errors(_invoke_api, operation="HEAD", bucket=bucket, key=key)
375
376 def _list_objects(
377 self,
378 path: str,
379 start_after: str | None = None,
380 end_at: str | None = None,
381 include_directories: bool = False,
382 symlink_handling: SymlinkHandling = SymlinkHandling.FOLLOW,
383 ) -> Iterator[ObjectMetadata]:
384 bucket, prefix = split_path(path)
385
386 # Get the prefix of the start_after and end_at paths relative to the bucket.
387 if start_after:
388 _, start_after = split_path(start_after)
389 if end_at:
390 _, end_at = split_path(end_at)
391
392 def _invoke_api() -> Iterator[ObjectMetadata]:
393 # AIS has no start key option like other object stores.
394 all_objects = self.client.bucket(bck_name=bucket, provider=self.provider).list_objects_iter(
395 prefix=prefix, props="name,size,atime,checksum,cone", page_size=DEFAULT_PAGE_SIZE
396 )
397
398 # Assume AIS guarantees lexicographical order.
399 for bucket_entry in all_objects:
400 obj = bucket_entry.object
401 key = obj.name
402 props = bucket_entry.generate_object_props()
403
404 # The access time is not always present in the response.
405 if props.access_time:
406 last_modified = dateutil_parser(props.access_time).astimezone(timezone.utc)
407 else:
408 last_modified = AWARE_DATETIME_MIN
409
410 if (start_after is None or start_after < key) and (end_at is None or key <= end_at):
411 yield ObjectMetadata(
412 key=key, content_length=int(props.size), last_modified=last_modified, etag=props.checksum_value
413 )
414 elif end_at is not None and end_at < key:
415 return
416
417 return self._translate_errors(_invoke_api, operation="LIST", bucket=bucket, key=prefix)
418
419 def _upload_file(self, remote_path: str, f: str | IO, attributes: dict[str, str] | None = None) -> int:
420 file_size: int = 0
421
422 if isinstance(f, str):
423 with open(f, "rb") as fp:
424 body = fp.read()
425 file_size = len(body)
426 self._put_object(remote_path, body, attributes=attributes)
427 else:
428 if isinstance(f, io.StringIO):
429 body = f.read().encode("utf-8")
430 file_size = len(body)
431 self._put_object(remote_path, body, attributes=attributes)
432 else:
433 body = f.read()
434 file_size = len(body)
435 self._put_object(remote_path, body, attributes=attributes)
436
437 return file_size
438
439 def _download_file(self, remote_path: str, f: str | IO, metadata: ObjectMetadata | None = None) -> int:
440 if metadata is None:
441 metadata = self._get_object_metadata(remote_path)
442
443 if isinstance(f, str):
444 if os.path.dirname(f):
445 safe_makedirs(os.path.dirname(f))
446 with open(f, "wb") as fp:
447 fp.write(self._get_object(remote_path))
448 else:
449 if isinstance(f, io.StringIO):
450 f.write(self._get_object(remote_path).decode("utf-8"))
451 else:
452 f.write(self._get_object(remote_path))
453
454 return metadata.content_length