-
Notifications
You must be signed in to change notification settings - Fork 189
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
sdk/python: Refactor internal object classes for accessing and iterat…
…ing over object content Signed-off-by: Aaron Wilson <[email protected]>
- Loading branch information
Showing
17 changed files
with
1,000 additions
and
610 deletions.
There are no files selected for viewing
Large diffs are not rendered by default.
Oops, something went wrong.
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,38 @@ | ||
# | ||
# Copyright (c) 2024, NVIDIA CORPORATION. All rights reserved. | ||
# | ||
from typing import Iterator | ||
|
||
from aistore.sdk.obj.object_client import ObjectClient | ||
|
||
|
||
# pylint: disable=too-few-public-methods | ||
class ContentIterator: | ||
""" | ||
Provide an iterator to open an HTTP response stream and read chunks of object content. | ||
Args: | ||
client (ObjectClient): Client for accessing contents of an individual object. | ||
chunk_size (int): Size of each chunk of data yielded from the response stream. | ||
""" | ||
|
||
def __init__(self, client: ObjectClient, chunk_size: int): | ||
self._client = client | ||
self._chunk_size = chunk_size | ||
|
||
def iter_from_position(self, start_position: int) -> Iterator[bytes]: | ||
""" | ||
Make a request to get a stream from the provided object starting at a specific byte position | ||
and yield chunks of the stream content. | ||
Args: | ||
start_position (int): The byte position from which to start reading. | ||
Returns: | ||
Iterator[bytes]: An iterator over each chunk of bytes in the object starting from the specific position | ||
""" | ||
stream = self._client.get(stream=True, start_position=start_position) | ||
try: | ||
yield from stream.iter_content(chunk_size=self._chunk_size) | ||
finally: | ||
stream.close() |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,72 @@ | ||
# | ||
# Copyright (c) 2024, NVIDIA CORPORATION. All rights reserved. | ||
# | ||
from typing import Optional, Dict, List | ||
|
||
import requests | ||
|
||
from aistore.sdk.const import HTTP_METHOD_GET, HTTP_METHOD_HEAD, HEADER_RANGE | ||
from aistore.sdk.obj.object_attributes import ObjectAttributes | ||
from aistore.sdk.request_client import RequestClient | ||
|
||
|
||
class ObjectClient: | ||
""" | ||
ObjectClient is a simple wrapper around a given RequestClient that makes requests to an individual object. | ||
Args: | ||
request_client (RequestClient): The RequestClient used to make HTTP requests | ||
path (str): URL Path to the object | ||
params (List[str]): Query parameters for the request | ||
headers (Optional[Dict[str, str]]): HTTP request headers | ||
""" | ||
|
||
def __init__( | ||
self, | ||
request_client: RequestClient, | ||
path: str, | ||
params: List[str], | ||
headers: Optional[Dict[str, str]] = None, | ||
): | ||
self._request_client = request_client | ||
self._request_path = path | ||
self._request_params = params | ||
self._request_headers = headers | ||
|
||
def get(self, stream: bool, start_position: int) -> requests.Response: | ||
""" | ||
Make a request to AIS to get the object content, optionally starting at a specific byte position. | ||
Args: | ||
stream (bool): If True, stream the response content. | ||
start_position (int): The byte position to start reading from. | ||
Returns: | ||
requests.Response: The response object from the request. | ||
""" | ||
headers = self._request_headers.copy() if self._request_headers else {} | ||
if start_position != 0: | ||
headers[HEADER_RANGE] = f"bytes={start_position}-" | ||
|
||
resp = self._request_client.request( | ||
HTTP_METHOD_GET, | ||
path=self._request_path, | ||
params=self._request_params, | ||
stream=stream, | ||
headers=headers, | ||
) | ||
resp.raise_for_status() | ||
return resp | ||
|
||
def head(self) -> ObjectAttributes: | ||
""" | ||
Make a head request to AIS to update and return only object attributes. | ||
Returns: | ||
`ObjectAttributes` containing metadata for this object. | ||
""" | ||
resp = self._request_client.request( | ||
HTTP_METHOD_HEAD, path=self._request_path, params=self._request_params | ||
) | ||
return ObjectAttributes(resp.headers) |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Oops, something went wrong.