sdk/python: Refactor module structure

Signed-off-by: Aaron Wilson <[email protected]>

python/aistore/pytorch/base_iter_dataset.py

@@ -8,7 +8,7 @@
 from typing import List, Union, Iterable, Dict, Iterator, Tuple
 import torch.utils.data as torch_utils
 from abc import ABC, abstractmethod
-from aistore.sdk.ais_source import AISSource
+from aistore.sdk import AISSource
 
 
 class AISBaseIterDataset(ABC, torch_utils.IterableDataset):

python/aistore/pytorch/base_map_dataset.py

@@ -7,8 +7,7 @@
 from typing import List, Union, Dict
 from abc import ABC, abstractmethod
 from torch.utils.data import Dataset
-from aistore.sdk.ais_source import AISSource
-from aistore.sdk.object import Object
+from aistore.sdk import AISSource, Object
 
 
 class AISBaseMapDataset(ABC, Dataset):

python/aistore/pytorch/iter_dataset.py

@@ -6,7 +6,7 @@
 
 from aistore.pytorch.base_iter_dataset import AISBaseIterDataset
 from typing import List, Union, Dict
-from aistore.sdk.ais_source import AISSource
+from aistore.sdk import AISSource
 from alive_progress import alive_it
 
 

python/aistore/pytorch/map_dataset.py

@@ -6,7 +6,7 @@
 
 from typing import List, Union, Dict
 
-from aistore.sdk.ais_source import AISSource
+from aistore.sdk import AISSource
 from aistore.pytorch.base_map_dataset import AISBaseMapDataset
 
 

python/aistore/pytorch/multishard_dataset.py

@@ -4,11 +4,9 @@
 Copyright (c) 2024, NVIDIA CORPORATION. All rights reserved.
 """
 
-from aistore.sdk.dataset.data_shard import DataShard
+from aistore.sdk import ArchiveConfig, DataShard, ListObjectFlag
 from aistore.sdk import Bucket
 from typing import Iterator, List, Iterable
-from aistore.sdk.list_object_flag import ListObjectFlag
-from aistore.sdk.types import ArchiveSettings
 from torch.utils.data import IterableDataset
 
 
@@ -62,5 +60,5 @@ def _get_shard_objects_iterator(
                     obj_name = obj.name.replace(f"{path}/", "", 1)
                     yield bucket.object(path).get(
                         etl_name=etl_name,
-                        archive_settings=ArchiveSettings(archpath=obj_name),
+                        archive_config=ArchiveConfig(archpath=obj_name),
                     ).read_all()

python/aistore/pytorch/shard_reader.py

@@ -6,14 +6,13 @@
 Copyright (c) 2024, NVIDIA CORPORATION. All rights reserved.
 """
 
-from aistore.sdk.bucket import Bucket
+from aistore.sdk import Bucket, ListObjectFlag
 from typing import Dict, Iterator, List, Union
 from aistore.pytorch.utils import get_basename, get_extension
 from aistore.pytorch.base_iter_dataset import AISBaseIterDataset
 from alive_progress import alive_it
 from io import BytesIO
 from tarfile import open, TarError
-from aistore.sdk.list_object_flag import ListObjectFlag
 
 
 class AISShardReader(AISBaseIterDataset):

python/aistore/sdk/__init__.py

@@ -1,5 +1,24 @@
+"""
+Import client-accessible components here to provide consistent imports via `from aistore.sdk import *`
+
+Copyright (c) 2021-2024, NVIDIA CORPORATION. All rights reserved.
+"""
+
+# Clients
 from aistore.sdk.client import Client
-from aistore.sdk.list_object_flag import ListObjectFlag
-from aistore.sdk.bucket import Bucket
-from aistore.sdk.namespace import Namespace
 from aistore.sdk.authn.authn_client import AuthNClient
+
+# Core components
+from aistore.sdk.cluster import Cluster
+from aistore.sdk.namespace import Namespace
+from aistore.sdk.ais_source import AISSource
+from aistore.sdk.bucket import Bucket
+from aistore.sdk.obj.object import Object
+from aistore.sdk.job import Job
+from aistore.sdk.etl.etl import Etl
+
+# Config objects, types and dataclasses
+from aistore.sdk.archive_config import ArchiveConfig
+from aistore.sdk.blob_download_config import BlobDownloadConfig
+from aistore.sdk.dataset.data_shard import DataShard
+from aistore.sdk.list_object_flag import ListObjectFlag

python/aistore/sdk/ais_source.py

@@ -3,7 +3,7 @@
 #
 from abc import ABC, abstractmethod
 from typing import Iterable
-from aistore.sdk.object import Object
+from aistore.sdk.obj.object import Object
 from aistore.sdk.request_client import RequestClient
 
 

python/aistore/sdk/archive_config.py

@@ -0,0 +1,58 @@
+#
+# Copyright (c) 2024, NVIDIA CORPORATION. All rights reserved.
+#
+from dataclasses import dataclass
+from enum import Enum
+
+
+class ArchiveMode(Enum):
+    """
+    Archive mode for getting files/objects from an archive in a bucket
+    See `MatchMode` enum in the cmn/archive/read.go
+    """
+
+    REGEXP = "regexp"
+    PREFIX = "prefix"
+    SUFFIX = "suffix"
+    SUBSTR = "substr"
+    WDSKEY = "wdskey"
+
+
+@dataclass
+class ArchiveConfig:
+    """
+    Configuration for extracting files from an archive
+
+    Attributes:
+    archpath (str, optional): If the object is an archive, use `archpath` to extract a single file
+        from the archive
+    regex (str, optional): A prefix, suffix, WebDataset key, or general-purpose regular expression
+        used to match filenames within the archive and select possibly multiple files
+    mode (ArchiveMode, optional): Specifies the mode of archive extraction when using `regex`
+
+    Example:
+        # Extract a single file from an archive
+        single_file_settings = ArchiveConfig(
+            archpath="path/to/your/file.txt"
+        )
+
+        # Extract multiple files from an archive
+        multi_file_settings = ArchiveConfig(
+            regex = "log",  # Retrieve all log files from the archive
+            mode=ArchiveMode.SUFFIX,
+        )
+    """
+
+    archpath: str = ""
+    regex: str = ""
+    mode: ArchiveMode = None
+
+    def __post_init__(self):
+        if self.mode and not self.regex:
+            raise ValueError("Archive mode requires archive regex")
+
+        if self.regex and not self.mode:
+            raise ValueError("Archive regex requires archive mode")
+
+        if self.regex and self.archpath:
+            raise ValueError("Cannot use both Archive regex and Archive path")

python/aistore/sdk/blob_download_config.py

@@ -0,0 +1,25 @@
+#
+# Copyright (c) 2024, NVIDIA CORPORATION. All rights reserved.
+#
+from dataclasses import dataclass
+
+
+@dataclass
+class BlobDownloadConfig:
+    """
+    Configuration for downloading objects using a blob downloader
+
+    Attributes:
+        chunk_size (str, optional): Chunk size for the blob downloader. It can be specified in IEC
+            or SI units, or as raw bytes (e.g., "4mb", "1MiB", "1048576", "128k")
+        num_workers (str, optional): Number of concurrent workers for the blob downloader
+
+    Example:
+        blob_settings = BlobDownloadConfig(
+            chunk_size="1MiB",  # 1 MiB per chunk
+            num_workers="5"     # 5 concurrent download workers
+        )
+    """
+
+    chunk_size: str = None
+    num_workers: str = None

python/aistore/sdk/bucket.py

@@ -13,8 +13,8 @@
 import requests
 
 from aistore.sdk.ais_source import AISSource
-from aistore.sdk.etl_const import DEFAULT_ETL_TIMEOUT
-from aistore.sdk.object_iterator import ObjectIterator
+from aistore.sdk.etl.etl_const import DEFAULT_ETL_TIMEOUT
+from aistore.sdk.obj.object_iterator import ObjectIterator
 from aistore.sdk.const import (
     ACT_COPY_BCK,
     ACT_CREATE_BCK,
@@ -58,7 +58,7 @@
 )
 from aistore.sdk.multiobj import ObjectGroup, ObjectRange
 from aistore.sdk.request_client import RequestClient
-from aistore.sdk.object import Object
+from aistore.sdk.obj.object import Object
 from aistore.sdk.types import (
     ActionMsg,
     BucketEntry,
@@ -73,7 +73,7 @@
 )
 from aistore.sdk.list_object_flag import ListObjectFlag
 from aistore.sdk.utils import validate_directory, get_file_size
-from aistore.sdk.object_props import ObjectProps
+from aistore.sdk.obj.object_props import ObjectProps
 
 Header = NewType("Header", requests.structures.CaseInsensitiveDict)
 

python/aistore/sdk/client.py

@@ -19,9 +19,9 @@
 from aistore.sdk.session_manager import SessionManager
 from aistore.sdk.types import Namespace
 from aistore.sdk.job import Job
-from aistore.sdk.etl import Etl
+from aistore.sdk.etl.etl import Etl
 from aistore.sdk.utils import parse_url
-from aistore.sdk.object import Object
+from aistore.sdk.obj.object import Object
 from aistore.sdk.errors import InvalidURLException
 
 

python/aistore/sdk/etl.py → python/aistore/sdk/etl/etl.py

@@ -17,7 +17,7 @@
     URL_PATH_ETL,
     UTF_ENCODING,
 )
-from aistore.sdk.etl_const import (
+from aistore.sdk.etl.etl_const import (
     ETL_SUPPORTED_PYTHON_VERSIONS,
     DEFAULT_ETL_RUNTIME,
     DEFAULT_ETL_COMM,

python/aistore/sdk/list_object_flag.py

@@ -1,3 +1,7 @@
+#
+# Copyright (c) 2023-2024, NVIDIA CORPORATION. All rights reserved.
+#
+
 from __future__ import annotations
 from enum import Enum
 

python/aistore/sdk/multiobj/object_group.py

@@ -16,8 +16,8 @@
     ACT_TRANSFORM_OBJECTS,
     ACT_ARCHIVE_OBJECTS,
 )
-from aistore.sdk.etl_const import DEFAULT_ETL_TIMEOUT
-from aistore.sdk.object import Object
+from aistore.sdk.etl.etl_const import DEFAULT_ETL_TIMEOUT
+from aistore.sdk.obj.object import Object
 from aistore.sdk.multiobj.object_names import ObjectNames
 from aistore.sdk.multiobj.object_range import ObjectRange
 from aistore.sdk.multiobj.object_template import ObjectTemplate