tostorchconnector 1.0.4__tar.gz → 1.0.5__tar.gz

{tostorchconnector-1.0.4/tostorchconnector.egg-info → tostorchconnector-1.0.5}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: tostorchconnector
-Version: 1.0.4
+Version: 1.0.5
 Summary: TOS connector integration for PyTorch
 Author-email: xiangshijian <xiangshijian@bytedance.com>
 Classifier: Development Status :: 4 - Beta
@@ -10,6 +10,7 @@ Classifier: Programming Language :: Python :: 3.9
 Classifier: Programming Language :: Python :: 3.10
 Classifier: Programming Language :: Python :: 3.11
 Classifier: Programming Language :: Python :: 3.12
+Classifier: Programming Language :: Python :: 3.13
 Classifier: License :: OSI Approved :: Apache Software License
 Classifier: Operating System :: OS Independent
 Classifier: Topic :: Utilities

{tostorchconnector-1.0.4 → tostorchconnector-1.0.5}/pyproject.toml

@@ -5,7 +5,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "tostorchconnector"
-version = "1.0.4"
+version = "1.0.5"
 description = "TOS connector integration for PyTorch"
 authors = [{ name = "xiangshijian", email = "xiangshijian@bytedance.com" }]
 requires-python = ">=3.8,<3.13"
@@ -18,6 +18,7 @@ classifiers = [
     "Programming Language :: Python :: 3.10",
     "Programming Language :: Python :: 3.11",
     "Programming Language :: Python :: 3.12",
+    "Programming Language :: Python :: 3.13",
     "License :: OSI Approved :: Apache Software License",
     "Operating System :: OS Independent",
     "Topic :: Utilities"

{tostorchconnector-1.0.4 → tostorchconnector-1.0.5}/tests/test_tos_dataset.py

@@ -2,7 +2,9 @@ import os
 import unittest
 
 from tostorchconnector import TosMapDataset, TosIterableDataset, TosCheckpoint
-from tostorchconnector.tos_client import CredentialProvider, TosLogConfig
+from tostorchconnector.tos_client import CredentialProvider, ReaderType
+
+USE_NATIVE_CLIENT = True
 
 
 class TestTosDataSet(unittest.TestCase):
@@ -14,7 +16,8 @@ class TestTosDataSet(unittest.TestCase):
         sk = os.getenv('TOS_SECRET_KEY')
         bucket = 'tos-pytorch-connector'
         datasets = TosMapDataset.from_urls(iter([f'tos://{bucket}/key1', f'tos://{bucket}/key2', f'{bucket}/key3']),
-                                           region=region, endpoint=endpoint, cred=CredentialProvider(ak, sk))
+                                           region=region, endpoint=endpoint, cred=CredentialProvider(ak, sk),
+                                           use_native_client=USE_NATIVE_CLIENT)
 
         for i in range(len(datasets)):
             print(datasets[i].bucket, datasets[i].key)
@@ -26,7 +29,8 @@ class TestTosDataSet(unittest.TestCase):
         sk = os.getenv('TOS_SECRET_KEY')
         bucket = 'tos-pytorch-connector'
         datasets = TosMapDataset.from_prefix(f'tos://{bucket}/prefix', region=region,
-                                             endpoint=endpoint, cred=CredentialProvider(ak, sk))
+                                             endpoint=endpoint, cred=CredentialProvider(ak, sk),
+                                             use_native_client=USE_NATIVE_CLIENT)
         for i in range(len(datasets)):
             item = datasets[i]
             print(item.bucket, item.key)
@@ -43,7 +47,8 @@ class TestTosDataSet(unittest.TestCase):
         sk = os.getenv('TOS_SECRET_KEY')
         bucket = 'tos-pytorch-connector'
         datasets = TosIterableDataset.from_prefix(f'tos://{bucket}/prefix', region=region,
-                                                  endpoint=endpoint, cred=CredentialProvider(ak, sk))
+                                                  endpoint=endpoint, cred=CredentialProvider(ak, sk),
+                                                  use_native_client=USE_NATIVE_CLIENT)
         i = 0
         for dataset in datasets:
             print(dataset.bucket, dataset.key)
@@ -59,15 +64,27 @@ class TestTosDataSet(unittest.TestCase):
         ak = os.getenv('TOS_ACCESS_KEY')
         sk = os.getenv('TOS_SECRET_KEY')
         bucket = 'tos-pytorch-connector'
-        checkpoint = TosCheckpoint(region, endpoint, cred=CredentialProvider(ak, sk), use_native_client=True)
+        checkpoint = TosCheckpoint(region, endpoint, cred=CredentialProvider(ak, sk),
+                                   use_native_client=USE_NATIVE_CLIENT)
         url = f'tos://{bucket}/key1'
+        print('test sequential')
         with checkpoint.writer(url) as writer:
             writer.write(b'hello world')
             writer.write(b'hi world')
 
-        reader = checkpoint.reader(url)
-        print(reader.read())
-
+        with checkpoint.reader(url) as reader:
+            data = reader.read(5)
+            print(data)
+            print(reader.read())
+            reader.seek(0)
+            data = reader.read(5)
+            print(data)
 
-if __name__ == '__main__':
-    unittest.main()
+        print('test ranged')
+        with checkpoint.reader(url, reader_type=ReaderType.RANGED) as reader:
+            data = reader.read(5)
+            print(data)
+            print(reader.read())
+            reader.seek(0)
+            data = reader.read(5)
+            print(data)

tostorchconnector-1.0.4/tests/test_tosnativeclient.py → tostorchconnector-1.0.5/tests/test_tosclient.py

@@ -5,7 +5,7 @@ import uuid
 from tosnativeclient import TosClient, TosException
 
 
-class TestTosNativeClient(unittest.TestCase):
+class TestTosClient(unittest.TestCase):
     def test_list_objects(self):
         region = os.getenv('TOS_REGION')
         endpoint = os.getenv('TOS_ENDPOINT')
@@ -15,16 +15,16 @@ class TestTosNativeClient(unittest.TestCase):
         tos_client = TosClient(region, endpoint, ak, sk, directives='info', directory='logs',
                                file_name_prefix='app.log')
 
-        list_stream = tos_client.list_objects(bucket, 'prefix/', max_keys=1000)
+        list_stream = tos_client.list_objects(bucket, '', max_keys=1000)
         count = 0
         try:
             for objects in list_stream:
-                count += 1
                 for content in objects.contents:
+                    count += 1
                     print(content.key, content.size)
-                    output = tos_client.head_object(bucket, content.key)
-                    assert output.etag == content.etag
-                    assert output.size == content.size
+                    # output = tos_client.head_object(bucket, content.key)
+                    # assert output.etag == content.etag
+                    # assert output.size == content.size
 
             print(count)
         except TosException as e:
@@ -70,7 +70,3 @@ class TestTosNativeClient(unittest.TestCase):
                 print(chunk)
         except TosException as e:
             print(e.args[0].status_code)
-
-
-if __name__ == '__main__':
-    unittest.main()

tostorchconnector-1.0.5/tests/test_tosrawclient.py

@@ -0,0 +1,102 @@
+import os
+import unittest
+import uuid
+
+import tos
+from tos.exceptions import TosServerError
+
+
+class TestTosRawClient(unittest.TestCase):
+    def test_put_get(self):
+        from tosnativeclient import TosRawClient, HeadObjectInput, TosException, DeleteObjectInput, \
+            PutObjectFromBufferInput, \
+            GetObjectInput, GetObjectOutput
+
+        region = os.getenv('TOS_REGION')
+        endpoint = os.getenv('TOS_ENDPOINT')
+        ak = os.getenv('TOS_ACCESS_KEY')
+        sk = os.getenv('TOS_SECRET_KEY')
+        bucket = 'tos-pytorch-connector'
+        key = str(uuid.uuid4())
+        tos_raw_client = TosRawClient(region, endpoint, ak, sk)
+
+        doutput = tos_raw_client.delete_object(DeleteObjectInput(bucket, key))
+        assert doutput.status_code == 204
+
+        try:
+            tos_raw_client.head_object(HeadObjectInput(bucket, key))
+            assert False
+        except TosException as e:
+            assert e.args[0].status_code == 404
+            assert len(e.args[0].request_id) > 0
+
+        data = str(uuid.uuid4()).encode('utf-8')
+        input = PutObjectFromBufferInput(bucket, key, data)
+        poutput = tos_raw_client.put_object_from_buffer(input)
+        assert poutput.status_code == 200
+        assert len(poutput.etag) > 0
+
+        houtput = tos_raw_client.head_object(HeadObjectInput(bucket, key))
+        assert houtput.status_code == 200
+        assert houtput.etag == poutput.etag
+        assert houtput.content_length == len(data)
+
+        goutput: GetObjectOutput = tos_raw_client.get_object(GetObjectInput(bucket, key))
+        assert goutput.status_code == 200
+        assert goutput.etag == poutput.etag
+        rdata = goutput.read_all()
+        assert rdata == data
+
+        doutput = tos_raw_client.delete_object(DeleteObjectInput(bucket, key))
+        assert doutput.status_code == 204
+
+        try:
+            tos_raw_client.get_object(GetObjectInput(bucket, key))
+            assert False
+        except TosException as e:
+            assert e.args[0].status_code == 404
+            assert len(e.args[0].request_id) > 0
+
+    def test_put_get_old(self):
+        region = os.getenv('TOS_REGION')
+        endpoint = os.getenv('TOS_ENDPOINT')
+        ak = os.getenv('TOS_ACCESS_KEY')
+        sk = os.getenv('TOS_SECRET_KEY')
+        bucket = 'tos-pytorch-connector'
+        key = str(uuid.uuid4())
+        tos_client = tos.TosClientV2(ak, sk, endpoint=endpoint, region=region)
+        doutput = tos_client.delete_object(bucket, key)
+        assert doutput.status_code == 204
+
+        try:
+            tos_client.head_object(bucket, key)
+            assert False
+        except TosServerError as e:
+            assert e.status_code == 404
+            assert len(e.request_id) > 0
+
+        data = str(uuid.uuid4()).encode('utf-8')
+        poutput = tos_client.put_object(bucket, key, content=data)
+        assert poutput.status_code == 200
+        assert len(poutput.etag) > 0
+
+        houtput = tos_client.head_object(bucket, key)
+        assert houtput.status_code == 200
+        assert houtput.etag == poutput.etag
+        assert houtput.content_length == len(data)
+
+        goutput = tos_client.get_object(bucket, key)
+        assert goutput.status_code == 200
+        assert goutput.etag == poutput.etag
+        rdata = goutput.read()
+        assert rdata == data
+
+        doutput = tos_client.delete_object(bucket, key)
+        assert doutput.status_code == 204
+
+        try:
+            tos_client.get_object(bucket, key)
+            assert False
+        except TosServerError as e:
+            assert e.status_code == 404
+            assert len(e.request_id) > 0

{tostorchconnector-1.0.4 → tostorchconnector-1.0.5}/tostorchconnector/__init__.py

@@ -1,4 +1,4 @@
-from .tos_object_reader import TosObjectReader
+from .tos_object_reader import TosObjectReader, SequentialTosObjectReader
 from .tos_object_writer import TosObjectWriter
 from .tos_iterable_dataset import TosIterableDataset
 from .tos_map_dataset import TosMapDataset
@@ -7,6 +7,7 @@ from tosnativeclient import TosException, TosError
 
 __all__ = [
     'TosObjectReader',
+    'SequentialTosObjectReader',
     'TosObjectWriter',
     'TosIterableDataset',
     'TosMapDataset',

{tostorchconnector-1.0.4 → tostorchconnector-1.0.5}/tostorchconnector/tos_checkpoint.py

@@ -2,7 +2,7 @@ import logging
 from typing import Optional
 
 from . import TosObjectReader, TosObjectWriter
-from .tos_client import CredentialProvider, TosClientConfig, TosClient, TosLogConfig
+from .tos_client import CredentialProvider, TosClientConfig, TosClient, TosLogConfig, ReaderType
 from .tos_common import parse_tos_url
 
 log = logging.getLogger(__name__)
@@ -14,26 +14,20 @@ class TosCheckpoint(object):
                  cred: Optional[CredentialProvider] = None,
                  client_conf: Optional[TosClientConfig] = None,
                  log_conf: Optional[TosLogConfig] = None, use_native_client=True):
-        self._client = None
-        self._native_client = None
         self._region = region
         self._endpoint = endpoint
         self._cred = cred
         self._client_conf = client_conf
         self._log_conf = log_conf
-        self._use_native_client = use_native_client
+        self._client = TosClient(self._region, self._endpoint, self._cred, self._client_conf, self._log_conf,
+                                 use_native_client)
+        log.info('TosCheckpoint init tos client succeed')
 
-    def reader(self, url: str) -> TosObjectReader:
+    def reader(self, url: str, reader_type: Optional[ReaderType] = None,
+               buffer_size: Optional[int] = None) -> TosObjectReader:
         bucket, key = parse_tos_url(url)
-        return self._get_tos_client().get_object(bucket, key)
+        return self._client.get_object(bucket, key, reader_type=reader_type, buffer_size=buffer_size)
 
     def writer(self, url: str) -> TosObjectWriter:
         bucket, key = parse_tos_url(url)
-        return self._get_tos_client().put_object(bucket, key)
-
-    def _get_tos_client(self):
-        if self._client is None:
-            self._client = TosClient(self._region, self._endpoint, self._cred, self._client_conf, self._log_conf,
-                                     self._use_native_client)
-            log.info('TosIterableDataset init tos client succeed')
-        return self._client
+        return self._client.put_object(bucket, key)

{tostorchconnector-1.0.4 → tostorchconnector-1.0.5}/tostorchconnector/tos_client.py

@@ -1,21 +1,26 @@
-import io
+import enum
 import logging
 import os
 from functools import partial
 from typing import Optional, List, Tuple
 
 import tos
-from tos.models2 import GetObjectOutput, PutObjectOutput
 
 import tosnativeclient
 
-from . import TosObjectReader, TosObjectWriter
+from . import SequentialTosObjectReader, TosObjectWriter, TosObjectReader
 from .tos_object_meta import TosObjectMeta
+from .tos_object_reader import TosObjectStream, RangedTosObjectReader
 from .tos_object_writer import PutObjectStream
 
 log = logging.getLogger(__name__)
 
 
+class ReaderType(enum.Enum):
+    SEQUENTIAL = 'Sequential'
+    RANGED = 'Ranged'
+
+
 class CredentialProvider(object):
     def __init__(self, ak: str, sk: str):
         self._ak = ak
@@ -110,48 +115,49 @@ class TosClient(object):
                 tos.set_logger(file_path=file_path, level=log_level)
 
     @property
-    def use_native_client(self):
+    def use_native_client(self) -> bool:
         return self._use_native_client
 
     def get_object(self, bucket: str, key: str, etag: Optional[str] = None,
-                   size: Optional[int] = None) -> TosObjectReader:
+                   size: Optional[int] = None, reader_type: Optional[ReaderType] = None,
+                   buffer_size: Optional[int] = None) -> TosObjectReader:
         log.debug(f'get_object tos://{bucket}/{key}')
+
         if size is None or etag is None:
             get_object_meta = partial(self.head_object, bucket, key)
         else:
             get_object_meta = lambda: TosObjectMeta(bucket, key, size, etag)
-        if isinstance(self._client, tosnativeclient.TosClient):
-            get_object_stream = partial(self._client.get_object, bucket, key)
-        else:
-            def get_object_stream(et: str, _: int) -> GetObjectOutput:
-                return self._client.get_object(bucket, key, '', et)
 
-        return TosObjectReader(bucket, key, get_object_meta, get_object_stream)
+        object_stream = TosObjectStream(bucket, key, get_object_meta, self._client)
+        if reader_type is not None and reader_type == ReaderType.RANGED:
+            return RangedTosObjectReader(bucket, key, object_stream, buffer_size)
+        return SequentialTosObjectReader(bucket, key, object_stream)
 
     def put_object(self, bucket: str, key: str, storage_class: Optional[str] = None) -> TosObjectWriter:
         log.debug(f'put_object tos://{bucket}/{key}')
 
         if isinstance(self._client, tosnativeclient.TosClient):
-            storage_class = storage_class if storage_class is not None else ''
-            put_object_stream = self._client.put_object(bucket, key, storage_class)
+            put_object_stream = self._client.put_object(bucket, key, storage_class=storage_class)
         else:
-            def put_object(content: io.BytesIO) -> PutObjectOutput:
-                return self._client.put_object(bucket, key, storage_class=storage_class, content=content)
-
-            put_object_stream = PutObjectStream(put_object)
+            put_object_stream = PutObjectStream(
+                lambda content: self._client.put_object(bucket, key, storage_class=storage_class, content=content))
 
         return TosObjectWriter(bucket, key, put_object_stream)
 
     def head_object(self, bucket: str, key: str) -> TosObjectMeta:
         log.debug(f'head_object tos://{bucket}/{key}')
+
         if isinstance(self._client, tosnativeclient.TosClient):
             resp = self._client.head_object(bucket, key)
             return TosObjectMeta(resp.bucket, resp.key, resp.size, resp.etag)
+
         resp = self._client.head_object(bucket, key)
         return TosObjectMeta(bucket, key, resp.content_length, resp.etag)
 
     def gen_list_stream(self, bucket: str, prefix: str, max_keys: int = 1000,
                         delimiter: Optional[str] = None) -> tosnativeclient.ListStream:
+        log.debug(f'gen_list_stream tos://{bucket}/{prefix}')
+
         if isinstance(self._client, tosnativeclient.TosClient):
             delimiter = delimiter if delimiter is not None else ''
             return self._client.list_objects(bucket, prefix, max_keys=max_keys, delimiter=delimiter)
@@ -161,6 +167,10 @@ class TosClient(object):
                      continuation_token: Optional[str] = None, delimiter: Optional[str] = None) -> Tuple[
         List[TosObjectMeta], bool, Optional[str]]:
         log.debug(f'list_objects tos://{bucket}/{prefix}')
+
+        if isinstance(self._client, tosnativeclient.TosClient):
+            raise NotImplementedError()
+
         resp = self._client.list_objects_type2(bucket, prefix, max_keys=max_keys, continuation_token=continuation_token,
                                                delimiter=delimiter)
         object_metas = []

{tostorchconnector-1.0.4 → tostorchconnector-1.0.5}/tostorchconnector/tos_common.py

@@ -22,14 +22,14 @@ class TosObjectIterator(object):
         self._is_truncated = True
         self._continuation_token = None
 
-    def close(self):
+    def close(self) -> None:
         if self._list_stream is not None:
             self._list_stream.close()
 
     def __iter__(self) -> Iterator[TosObjectMeta]:
         return self
 
-    def __next__(self):
+    def __next__(self) -> TosObjectMeta:
         if self._client.use_native_client:
             if self._list_stream is None:
                 self._list_stream = self._client.gen_list_stream(self._bucket, self._prefix, max_keys=1000,
@@ -37,9 +37,14 @@ class TosObjectIterator(object):
 
             if self._object_metas is None or self._index >= len(self._object_metas):
                 self._object_metas = []
-                objects = next(self._list_stream)
-                for content in objects.contents:
-                    self._object_metas.append(TosObjectMeta(content.bucket, content.key, content.size, content.etag))
+                self._index = 0
+                while 1:
+                    objects = next(self._list_stream)
+                    for content in objects.contents:
+                        self._object_metas.append(
+                            TosObjectMeta(content.bucket, content.key, content.size, content.etag))
+                    if len(self._object_metas) > 0:
+                        break
 
             object_meta = self._object_metas[self._index]
             self._index += 1

{tostorchconnector-1.0.4 → tostorchconnector-1.0.5}/tostorchconnector/tos_iterable_dataset.py

@@ -5,7 +5,7 @@ from typing import Iterator, Any, Optional, Callable, Union
 import torch
 
 from . import TosObjectReader
-from .tos_client import CredentialProvider, TosClientConfig, TosClient, TosLogConfig
+from .tos_client import CredentialProvider, TosClientConfig, TosClient, TosLogConfig, ReaderType
 from .tos_common import default_trans, gen_dataset_from_urls, gen_dataset_from_prefix
 from .tos_object_meta import TosObjectMeta
 
@@ -20,8 +20,10 @@ class TosIterableDataset(torch.utils.data.IterableDataset):
                  cred: Optional[CredentialProvider] = None,
                  client_conf: Optional[TosClientConfig] = None,
                  log_conf: Optional[TosLogConfig] = None,
-                 sharding: bool = False, use_native_client=True):
-        self._client: Optional[TosClient] = None
+                 sharding: bool = False,
+                 use_native_client: bool = True,
+                 reader_type: Optional[ReaderType] = None,
+                 buffer_size: Optional[int] = None):
         self._gen_dataset = gen_dataset
         self._region = region
         self._endpoint = endpoint
@@ -30,13 +32,17 @@ class TosIterableDataset(torch.utils.data.IterableDataset):
         self._client_conf = client_conf
         self._log_conf = log_conf
         self._sharding = sharding
-        self._use_native_client = use_native_client
         if torch.distributed.is_initialized():
             self._rank = torch.distributed.get_rank()
             self._world_size = torch.distributed.get_world_size()
         else:
             self._rank = 0
             self._world_size = 1
+        self._reader_type = reader_type
+        self._buffer_size = buffer_size
+        self._client = TosClient(self._region, self._endpoint, self._cred, self._client_conf, self._log_conf,
+                                 use_native_client)
+        log.info('TosIterableDataset init tos client succeed')
 
     @classmethod
     def from_urls(cls, urls: Union[str, Iterator[str]], *, region: str, endpoint: Optional[str] = None,
@@ -44,10 +50,13 @@ class TosIterableDataset(torch.utils.data.IterableDataset):
                   cred: Optional[CredentialProvider] = None,
                   client_conf: Optional[TosClientConfig] = None,
                   log_conf: Optional[TosLogConfig] = None,
-                  sharding: bool = False, use_native_client=True):
+                  sharding: bool = False,
+                  use_native_client: bool = True,
+                  reader_type: Optional[ReaderType] = None,
+                  buffer_size: Optional[int] = None):
         log.info(f'building {cls.__name__} from_urls')
         return cls(region, partial(gen_dataset_from_urls, urls), endpoint, transform, cred, client_conf, log_conf,
-                   sharding, use_native_client)
+                   sharding, use_native_client, reader_type, buffer_size)
 
     @classmethod
     def from_prefix(cls, prefix: str, *, region: str, endpoint: Optional[str] = None,
@@ -55,10 +64,13 @@ class TosIterableDataset(torch.utils.data.IterableDataset):
                     cred: Optional[CredentialProvider] = None,
                     client_conf: Optional[TosClientConfig] = None,
                     log_conf: Optional[TosLogConfig] = None,
-                    sharding: bool = False, use_native_client=True):
+                    sharding: bool = False,
+                    use_native_client: bool = True,
+                    reader_type: Optional[ReaderType] = None,
+                    buffer_size: Optional[int] = None):
         log.info(f'building {cls.__name__} from_prefix')
         return cls(region, partial(gen_dataset_from_prefix, prefix), endpoint, transform, cred, client_conf, log_conf,
-                   sharding, use_native_client)
+                   sharding, use_native_client, reader_type, buffer_size)
 
     def __iter__(self) -> Iterator[Any]:
         worker_id = 0
@@ -72,12 +84,12 @@ class TosIterableDataset(torch.utils.data.IterableDataset):
         if not self._sharding or (self._world_size == 1 and num_workers == 1):
             return map(
                 self._trans_tos_object,
-                self._gen_dataset(self._get_tos_client()),
+                self._gen_dataset(self._client),
             )
 
         part_dataset = (
             obj
-            for idx, obj in enumerate(self._gen_dataset(self._get_tos_client()))
+            for idx, obj in enumerate(self._gen_dataset(self._client))
             if idx % self._world_size == self._rank
         )
 
@@ -88,12 +100,7 @@ class TosIterableDataset(torch.utils.data.IterableDataset):
         )
         return map(self._trans_tos_object, part_dataset)
 
-    def _get_tos_client(self):
-        if self._client is None:
-            self._client = TosClient(self._region, self._endpoint, self._cred, self._client_conf, self._log_conf)
-            log.info('TosIterableDataset init tos client succeed')
-        return self._client
-
     def _trans_tos_object(self, object_meta: TosObjectMeta) -> Any:
-        obj = self._get_tos_client().get_object(object_meta.bucket, object_meta.key, object_meta.etag, object_meta.size)
+        obj = self._client.get_object(object_meta.bucket, object_meta.key, object_meta.etag, object_meta.size,
+                                      reader_type=self._reader_type, buffer_size=self._buffer_size)
         return self._trans(obj)

{tostorchconnector-1.0.4 → tostorchconnector-1.0.5}/tostorchconnector/tos_map_dataset.py

@@ -5,7 +5,7 @@ from typing import Any, Callable, Iterator, Optional, List, Union
 import torch
 
 from . import TosObjectReader
-from .tos_client import CredentialProvider, TosClientConfig, TosClient, TosLogConfig
+from .tos_client import CredentialProvider, TosClientConfig, TosClient, TosLogConfig, ReaderType
 from .tos_common import default_trans, gen_dataset_from_prefix, \
     gen_dataset_from_urls
 from .tos_object_meta import TosObjectMeta
@@ -20,8 +20,10 @@ class TosMapDataset(torch.utils.data.Dataset):
                  trans: Callable[[TosObjectReader], Any] = default_trans,
                  cred: Optional[CredentialProvider] = None,
                  client_conf: Optional[TosClientConfig] = None,
-                 log_conf: Optional[TosLogConfig] = None, use_native_client=True):
-        self._client: Optional[TosClient] = None
+                 log_conf: Optional[TosLogConfig] = None,
+                 use_native_client: bool = True,
+                 reader_type: Optional[ReaderType] = None,
+                 buffer_size: Optional[int] = None):
         self._gen_dataset = gen_dataset
         self._region = region
         self._endpoint = endpoint
@@ -29,28 +31,38 @@ class TosMapDataset(torch.utils.data.Dataset):
         self._cred = cred
         self._client_conf = client_conf
         self._log_conf = log_conf
-        self._use_native_client = use_native_client
         self._dataset: Optional[List[TosObjectMeta]] = None
+        self._reader_type = reader_type
+        self._buffer_size = buffer_size
+        self._client = TosClient(self._region, self._endpoint, self._cred, self._client_conf, self._log_conf,
+                                 use_native_client)
+        log.info('TosMapDataset init tos client succeed')
 
     @classmethod
     def from_urls(cls, urls: Union[str, Iterator[str]], *, region: str, endpoint: Optional[str] = None,
                   trans: Callable[[TosObjectReader], Any] = default_trans,
                   cred: Optional[CredentialProvider] = None,
                   client_conf: Optional[TosClientConfig] = None,
-                  log_conf: Optional[TosLogConfig] = None, use_native_client=True):
+                  log_conf: Optional[TosLogConfig] = None,
+                  use_native_client: bool = True,
+                  reader_type: Optional[ReaderType] = None,
+                  buffer_size: Optional[int] = None):
         log.info(f'building {cls.__name__} from_urls')
         return cls(region, partial(gen_dataset_from_urls, urls), endpoint, trans, cred, client_conf, log_conf,
-                   use_native_client)
+                   use_native_client, reader_type, buffer_size)
 
     @classmethod
     def from_prefix(cls, prefix: str, *, region: str, endpoint: Optional[str] = None,
                     trans: Callable[[TosObjectReader], Any] = default_trans,
                     cred: Optional[CredentialProvider] = None,
                     client_conf: Optional[TosClientConfig] = None,
-                    log_conf: Optional[TosLogConfig] = None, use_native_client=True):
+                    log_conf: Optional[TosLogConfig] = None,
+                    use_native_client: bool = True,
+                    reader_type: Optional[ReaderType] = None,
+                    buffer_size: Optional[int] = None):
         log.info(f'building {cls.__name__} from_prefix')
         return cls(region, partial(gen_dataset_from_prefix, prefix), endpoint, trans, cred, client_conf, log_conf,
-                   use_native_client)
+                   use_native_client, reader_type, buffer_size)
 
     def __getitem__(self, i: int) -> Any:
         return self._trans_tos_object(i)
@@ -58,21 +70,15 @@ class TosMapDataset(torch.utils.data.Dataset):
     def __len__(self) -> int:
         return len(self._data_set)
 
-    def _get_tos_client(self) -> TosClient:
-        if self._client is None:
-            self._client = TosClient(self._region, self._endpoint, self._cred, self._client_conf, self._log_conf,
-                                     self._use_native_client)
-            log.info('TosMapDataset init tos client succeed')
-        return self._client
-
     @property
-    def _data_set(self):
+    def _data_set(self) -> List[TosObjectMeta]:
         if self._dataset is None:
-            self._dataset = list(self._gen_dataset(self._get_tos_client()))
+            self._dataset = list(self._gen_dataset(self._client))
         assert self._dataset is not None
         return self._dataset
 
     def _trans_tos_object(self, i: int) -> Any:
         object_meta = self._data_set[i]
-        obj = self._get_tos_client().get_object(object_meta.bucket, object_meta.key, object_meta.etag, object_meta.size)
+        obj = self._client.get_object(object_meta.bucket, object_meta.key, object_meta.etag, object_meta.size,
+                                      reader_type=self._reader_type, buffer_size=self._buffer_size)
         return self._trans(obj)