tostorchconnector 1.0.6__tar.gz → 1.1.2__tar.gz

{tostorchconnector-1.0.6/tostorchconnector.egg-info → tostorchconnector-1.1.2}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: tostorchconnector
-Version: 1.0.6
+Version: 1.1.2
 Summary: TOS connector integration for PyTorch
 Author-email: xiangshijian <xiangshijian@bytedance.com>
 Classifier: Development Status :: 4 - Beta
@@ -19,7 +19,7 @@ Description-Content-Type: text/markdown
 License-File: LICENSE
 Requires-Dist: torch>=2.0
 Requires-Dist: tos>=2.8.0
-Requires-Dist: tosnativeclient>=1.0.2
+Requires-Dist: tosnativeclient>=1.1.1
 Dynamic: license-file
 
 # TOS Connector for pytorch

{tostorchconnector-1.0.6 → tostorchconnector-1.1.2}/pyproject.toml

@@ -5,7 +5,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "tostorchconnector"
-version = "1.0.6"
+version = "1.1.2"
 description = "TOS connector integration for PyTorch"
 authors = [{ name = "xiangshijian", email = "xiangshijian@bytedance.com" }]
 requires-python = ">=3.8,<3.14"
@@ -26,7 +26,7 @@ classifiers = [
 dependencies = [
     "torch >= 2.0",
     "tos>=2.8.0",
-    "tosnativeclient >= 1.0.2"
+    "tosnativeclient >= 1.1.1"
 ]
 
 [tool.setuptools.packages]

{tostorchconnector-1.0.6 → tostorchconnector-1.1.2}/tests/test_tos_dataset.py

@@ -1,10 +1,12 @@
 import os
+import pickle
 import unittest
 
 from tostorchconnector import TosMapDataset, TosIterableDataset, TosCheckpoint
 from tostorchconnector.tos_client import CredentialProvider, ReaderType
 
 USE_NATIVE_CLIENT = True
+READER_TYPE = ReaderType.SEQUENTIAL
 
 
 class TestTosDataSet(unittest.TestCase):
@@ -22,23 +24,50 @@ class TestTosDataSet(unittest.TestCase):
         for i in range(len(datasets)):
             print(datasets[i].bucket, datasets[i].key)
 
+    def test_pickle(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'
+        datasets = TosMapDataset.from_prefix(f'tos://{bucket}', region=region,
+                                             endpoint=endpoint, cred=CredentialProvider(ak, sk),
+                                             use_native_client=USE_NATIVE_CLIENT)
+        pickled_datasets = pickle.dumps(datasets)
+        assert isinstance(pickled_datasets, bytes)
+        unpickled_datasets = pickle.loads(pickled_datasets)
+        i = 0
+        for dataset in unpickled_datasets:
+            print(dataset.bucket, dataset.key)
+            i += 1
+        print(i)
+
+        pickled = pickle.dumps(unpickled_datasets._data_set)
+        assert isinstance(pickled, bytes)
+
     def test_from_prefix(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'
-        datasets = TosMapDataset.from_prefix(f'tos://{bucket}/prefix', region=region,
+        datasets = TosMapDataset.from_prefix(f'tos://{bucket}', region=region,
                                              endpoint=endpoint, cred=CredentialProvider(ak, sk),
                                              use_native_client=USE_NATIVE_CLIENT)
+
+        count = 0
         for i in range(len(datasets)):
-            item = datasets[i]
-            print(item.bucket, item.key)
-            if i == 1:
-                item = datasets[i]
-                data = item.read(100)
-                print(data)
-                print(len(data))
+            dataset = datasets[i]
+            print(dataset.bucket, dataset.key)
+            count += 1
+            dataset.close()
+            # if i == 1:
+            #     item = datasets[i]
+            #     data = item.read(100)
+            #     print(data)
+            #     print(len(data))
+
+        print(count)
 
     def test_from_prefix_iter(self):
         region = os.getenv('TOS_REGION')
@@ -46,17 +75,23 @@ class TestTosDataSet(unittest.TestCase):
         ak = os.getenv('TOS_ACCESS_KEY')
         sk = os.getenv('TOS_SECRET_KEY')
         bucket = 'tos-pytorch-connector'
-        datasets = TosIterableDataset.from_prefix(f'tos://{bucket}/prefix', region=region,
+        datasets = TosIterableDataset.from_prefix(f'tos://{bucket}', region=region,
                                                   endpoint=endpoint, cred=CredentialProvider(ak, sk),
-                                                  use_native_client=USE_NATIVE_CLIENT)
+                                                  use_native_client=USE_NATIVE_CLIENT, reader_type=ReaderType.RANGED)
         i = 0
         for dataset in datasets:
             print(dataset.bucket, dataset.key)
-            if i == 1:
-                data = dataset.read(100)
-                print(data)
-                print(len(data))
+            # if dataset.key == 'tosutil':
+            #     with open('logs/tosutil', 'wb') as f:
+            #         while 1:
+            #             chunk = dataset.read(8192)
+            #             print(len(chunk))
+            #             if not chunk:
+            #                 break
+            #             f.write(chunk)
+            dataset.close()
             i += 1
+        print(i)
 
     def test_checkpoint(self):
         region = os.getenv('TOS_REGION')
@@ -72,6 +107,9 @@ class TestTosDataSet(unittest.TestCase):
             writer.write(b'hello world')
             writer.write(b'hi world')
 
+        with checkpoint.reader(url) as reader:
+            print(reader.read())
+
         with checkpoint.reader(url) as reader:
             data = reader.read(5)
             print(data)

{tostorchconnector-1.0.6 → tostorchconnector-1.1.2}/tests/test_tosclient.py

@@ -1,4 +1,5 @@
 import os
+import pickle
 import unittest
 import uuid
 
@@ -6,15 +7,40 @@ from tosnativeclient import TosClient, TosException
 
 
 class TestTosClient(unittest.TestCase):
+    def test_pickle(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'
+        tos_client = TosClient(region, endpoint, ak, sk)
+        pickled_tos_client = pickle.dumps(tos_client)
+        assert isinstance(pickled_tos_client, bytes)
+        unpickled_tos_client = pickle.loads(pickled_tos_client)
+
+        self._test_list_objects(bucket, unpickled_tos_client)
+        self._test_write_read_object(bucket, unpickled_tos_client)
+
     def test_list_objects(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'
-        tos_client = TosClient(region, endpoint, ak, sk, directives='info', directory='logs',
-                               file_name_prefix='app.log')
+        tos_client = TosClient(region, endpoint, ak, sk)
+        self._test_list_objects(bucket, tos_client)
 
+    def test_write_read_object(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'
+        tos_client = TosClient(region, endpoint, ak, sk)
+
+        self._test_write_read_object(bucket, tos_client)
+
+    def _test_list_objects(self, bucket, tos_client):
         list_stream = tos_client.list_objects(bucket, '', max_keys=1000)
         count = 0
         try:
@@ -22,23 +48,15 @@ class TestTosClient(unittest.TestCase):
                 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:
             print(e.args[0].message)
 
-    def test_write_read_object(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'
-        tos_client = TosClient(region, endpoint, ak, sk, directives='info', directory='logs',
-                               file_name_prefix='app.log')
-
+    def _test_write_read_object(self, bucket, tos_client):
         key = str(uuid.uuid4())
         read_stream = tos_client.get_object(bucket, key, '', 1)
 
@@ -59,6 +77,8 @@ class TestTosClient(unittest.TestCase):
         write_stream.close()
 
         output = tos_client.head_object(bucket, key)
+        print(output.etag, output.size)
+
         read_stream = tos_client.get_object(bucket, key, output.etag, output.size)
         try:
             offset = 0

{tostorchconnector-1.0.6 → tostorchconnector-1.1.2}/tostorchconnector/tos_checkpoint.py

@@ -13,14 +13,14 @@ class TosCheckpoint(object):
                  endpoint: Optional[str] = None,
                  cred: Optional[CredentialProvider] = None,
                  client_conf: Optional[TosClientConfig] = None,
-                 log_conf: Optional[TosLogConfig] = None, use_native_client=True):
+                 use_native_client=True,
+                 enable_crc=True):
         self._region = region
         self._endpoint = endpoint
         self._cred = cred
         self._client_conf = client_conf
-        self._log_conf = log_conf
-        self._client = TosClient(self._region, self._endpoint, self._cred, self._client_conf, self._log_conf,
-                                 use_native_client)
+        self._client = TosClient(self._region, self._endpoint, self._cred, self._client_conf, use_native_client,
+                                 enable_crc)
         log.info('TosCheckpoint init tos client succeed')
 
     def reader(self, url: str, reader_type: Optional[ReaderType] = None,

{tostorchconnector-1.0.6 → tostorchconnector-1.1.2}/tostorchconnector/tos_client.py

@@ -1,8 +1,10 @@
 import enum
 import logging
 import os
+import gc
+
 from functools import partial
-from typing import Optional, List, Tuple
+from typing import Optional, List, Tuple, Any
 
 import tos
 
@@ -15,6 +17,44 @@ from .tos_object_writer import PutObjectStream
 
 log = logging.getLogger(__name__)
 
+import threading
+import weakref
+import traceback
+
+_client_lock = threading.Lock()
+_client_map = weakref.WeakSet()
+
+
+def _before_fork():
+    with _client_lock:
+        clients = list(_client_map)
+
+    if not clients or len(clients) == 0:
+        return
+
+    try:
+        for client in clients:
+            client._inner_client = None
+
+        _reset_client_map()
+        gc.collect()
+    except Exception as e:
+        log.warning(f'failed to clean up native clients, {str(e)}')
+        traceback.print_exc()
+
+
+def _after_fork_in_child():
+    _reset_client_map()
+
+
+def _reset_client_map():
+    global _client_map
+    with _client_lock:
+        _client_map = weakref.WeakSet()
+
+
+os.register_at_fork(before=_before_fork, after_in_child=_after_fork_in_child)
+
 
 class ReaderType(enum.Enum):
     SEQUENTIAL = 'Sequential'
@@ -77,42 +117,38 @@ class TosLogConfig(object):
 
 class TosClient(object):
     def __init__(self, region: str, endpoint: Optional[str] = None, cred: Optional[CredentialProvider] = None,
-                 client_conf: Optional[TosClientConfig] = None, log_conf: Optional[TosLogConfig] = None,
-                 use_native_client: bool = True):
-        cred = CredentialProvider('', '') if cred is None else cred
-        client_conf = TosClientConfig() if client_conf is None else client_conf
-        log_conf = TosLogConfig() if log_conf is None else log_conf
-        self._part_size = client_conf.part_size
+                 client_conf: Optional[TosClientConfig] = None, use_native_client: bool = True,
+                 enable_crc: bool = True):
+        self._region = region
+        self._endpoint = endpoint
+        self._cred = CredentialProvider('', '') if cred is None else cred
+        self._client_conf = TosClientConfig() if client_conf is None else client_conf
+        self._part_size = self._client_conf.part_size
         self._use_native_client = use_native_client
-        if use_native_client:
-            directives = ''
-            directory = ''
-            file_name_prefix = ''
-            if log_conf.log_dir and log_conf.log_file_name:
-                if log_conf.log_level:
-                    if log_conf.log_level == logging.DEBUG:
-                        directives = 'debug'
-                    elif log_conf.log_level == logging.INFO:
-                        directives = 'info'
-                    elif log_conf.log_level == logging.WARN:
-                        directives = 'warn'
-                    elif log_conf.log_level == logging.ERROR:
-                        directives = 'error'
-                    else:
-                        directives = 'info'
-                directory = log_conf.log_dir
-                file_name_prefix = log_conf.log_file_name
-            self._client = tosnativeclient.TosClient(region, endpoint, cred.ak, cred.sk, client_conf.part_size,
-                                                     client_conf.max_retry_count, directives=directives,
-                                                     directory=directory,
-                                                     file_name_prefix=file_name_prefix)
-        else:
-            self._client = tos.TosClientV2(cred.ak, cred.sk, endpoint=endpoint, region=region,
-                                           max_retry_count=client_conf.max_retry_count)
-            if log_conf.log_dir and log_conf.log_file_name:
-                file_path = os.path.join(log_conf.log_dir, log_conf.log_file_name)
-                log_level = log_conf.log_level if log_conf.log_level else logging.INFO
-                tos.set_logger(file_path=file_path, level=log_level)
+        self._inner_client = None
+        self._client_pid = None
+        self._enable_crc = enable_crc
+
+    @property
+    def _client(self) -> Any:
+        if self._client_pid is None or self._client_pid != os.getpid() or self._inner_client is None:
+            with _client_lock:
+                if self._use_native_client:
+                    self._inner_client = tosnativeclient.TosClient(self._region, self._endpoint, self._cred.ak,
+                                                                   self._cred.sk,
+                                                                   self._client_conf.part_size,
+                                                                   self._client_conf.max_retry_count,
+                                                                   enable_crc=self._enable_crc)
+                else:
+                    self._inner_client = tos.TosClientV2(self._cred.ak, self._cred.sk, endpoint=self._endpoint,
+                                                         region=self._region,
+                                                         max_retry_count=self._client_conf.max_retry_count,
+                                                         enable_crc=self._enable_crc)
+                self._client_pid = os.getpid()
+                _client_map.add(self)
+
+        assert self._inner_client is not None
+        return self._inner_client
 
     @property
     def use_native_client(self) -> bool:
@@ -155,12 +191,17 @@ class TosClient(object):
         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:
+                        delimiter: Optional[str] = None,
+                        continuation_token: Optional[str] = None,
+                        list_background_buffer_count: int = 1) -> 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)
+            continuation_token = continuation_token if continuation_token is not None else ''
+            return self._client.list_objects(bucket, prefix, max_keys=max_keys, delimiter=delimiter,
+                                             continuation_token=continuation_token,
+                                             list_background_buffer_count=list_background_buffer_count)
         raise NotImplementedError()
 
     def list_objects(self, bucket: str, prefix: str, max_keys: int = 1000,

{tostorchconnector-1.0.6 → tostorchconnector-1.1.2}/tostorchconnector/tos_common.py

@@ -1,6 +1,7 @@
 import logging
 from typing import Union, Iterator, Tuple, Optional
 
+from tosnativeclient import TosObject
 from . import TosObjectReader
 from .tos_client import TosClient
 from .tos_object_meta import TosObjectMeta
@@ -8,19 +9,30 @@ from .tos_object_meta import TosObjectMeta
 log = logging.getLogger(__name__)
 
 
-class TosObjectIterator(object):
+class TosObjectIterable(object):
     def __init__(self, bucket: str, prefix: str, client: TosClient):
         self._bucket = bucket
         self._prefix = prefix
+        self._list_background_buffer_count = 3
+        self._client = client
+
+    def __iter__(self) -> Iterator[TosObjectMeta]:
+        return iter(TosObjectIterator(self._bucket, self._prefix, self._list_background_buffer_count, self._client))
+
+
+class TosObjectIterator(object):
+    def __init__(self, bucket: str, prefix: str, list_background_buffer_count: int, client: TosClient):
+        self._bucket = bucket
+        self._prefix = prefix
+        self._list_background_buffer_count = list_background_buffer_count
         self._client = client
         self._delimiter: Optional[str] = None
-        self._list_stream = None
+        self._continuation_token: Optional[str] = None
 
+        self._list_stream = None
         self._object_metas = None
         self._index = 0
-
         self._is_truncated = True
-        self._continuation_token = None
 
     def close(self) -> None:
         if self._list_stream is not None:
@@ -33,22 +45,28 @@ class TosObjectIterator(object):
         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,
-                                                                 delimiter=self._delimiter)
+                                                                 delimiter=self._delimiter,
+                                                                 continuation_token=self._continuation_token,
+                                                                 list_background_buffer_count=self._list_background_buffer_count)
 
             if self._object_metas is None or self._index >= len(self._object_metas):
-                self._object_metas = []
+                self._object_metas = None
                 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:
+                    try:
+                        objects = next(self._list_stream)
+                    except:
+                        self.close()
+                        raise
+                    self._continuation_token = self._list_stream.current_continuation_token()
+                    self._object_metas = objects.contents
+                    if self._object_metas is not None and len(self._object_metas) > 0:
                         break
 
             object_meta = self._object_metas[self._index]
             self._index += 1
-            return object_meta
+            # this is very critical
+            return TosObjectMeta(self._bucket, object_meta.key, object_meta.size, object_meta.etag)
 
         while self._object_metas is None or self._index >= len(self._object_metas):
             if not self._is_truncated:
@@ -101,4 +119,4 @@ def gen_dataset_from_urls(urls: Union[str, Iterator[str]], _: TosClient) -> Iter
 
 def gen_dataset_from_prefix(prefix: str, client: TosClient) -> Iterator[TosObjectMeta]:
     bucket, prefix = parse_tos_url(prefix)
-    return iter(TosObjectIterator(bucket, prefix, client))
+    return iter(TosObjectIterable(bucket, prefix, client))

{tostorchconnector-1.0.6 → tostorchconnector-1.1.2}/tostorchconnector/tos_iterable_dataset.py

@@ -4,8 +4,9 @@ from typing import Iterator, Any, Optional, Callable, Union
 
 import torch
 
+from tosnativeclient import TosObject
 from . import TosObjectReader
-from .tos_client import CredentialProvider, TosClientConfig, TosClient, TosLogConfig, ReaderType
+from .tos_client import CredentialProvider, TosClientConfig, TosClient, ReaderType
 from .tos_common import default_trans, gen_dataset_from_urls, gen_dataset_from_prefix
 from .tos_object_meta import TosObjectMeta
 
@@ -16,21 +17,20 @@ class TosIterableDataset(torch.utils.data.IterableDataset):
     def __init__(self, region: str,
                  gen_dataset: Callable[[TosClient], Iterator[TosObjectMeta]],
                  endpoint: Optional[str] = None,
-                 trans: Callable[[TosObjectReader], Any] = default_trans,
+                 transform: Callable[[TosObjectReader], Any] = default_trans,
                  cred: Optional[CredentialProvider] = None,
                  client_conf: Optional[TosClientConfig] = None,
-                 log_conf: Optional[TosLogConfig] = None,
                  sharding: bool = False,
                  use_native_client: bool = True,
                  reader_type: Optional[ReaderType] = None,
-                 buffer_size: Optional[int] = None):
+                 buffer_size: Optional[int] = None,
+                 enable_crc: bool = True):
         self._gen_dataset = gen_dataset
         self._region = region
         self._endpoint = endpoint
-        self._trans = trans
+        self._trans = transform
         self._cred = cred
         self._client_conf = client_conf
-        self._log_conf = log_conf
         self._sharding = sharding
         if torch.distributed.is_initialized():
             self._rank = torch.distributed.get_rank()
@@ -40,8 +40,8 @@ class TosIterableDataset(torch.utils.data.IterableDataset):
             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)
+        self._client = TosClient(self._region, self._endpoint, self._cred, self._client_conf,
+                                 use_native_client, enable_crc)
         log.info('TosIterableDataset init tos client succeed')
 
     @classmethod
@@ -49,28 +49,28 @@ class TosIterableDataset(torch.utils.data.IterableDataset):
                   transform: Callable[[TosObjectReader], Any] = default_trans,
                   cred: Optional[CredentialProvider] = None,
                   client_conf: Optional[TosClientConfig] = None,
-                  log_conf: Optional[TosLogConfig] = None,
                   sharding: bool = False,
                   use_native_client: bool = True,
                   reader_type: Optional[ReaderType] = None,
-                  buffer_size: Optional[int] = None):
+                  buffer_size: Optional[int] = None,
+                  enable_crc: bool = True):
         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, reader_type, buffer_size)
+        return cls(region, partial(gen_dataset_from_urls, urls), endpoint, transform, cred, client_conf,
+                   sharding, use_native_client, reader_type, buffer_size, enable_crc)
 
     @classmethod
     def from_prefix(cls, prefix: str, *, region: str, endpoint: Optional[str] = None,
                     transform: Callable[[TosObjectReader], Any] = default_trans,
                     cred: Optional[CredentialProvider] = None,
                     client_conf: Optional[TosClientConfig] = None,
-                    log_conf: Optional[TosLogConfig] = None,
                     sharding: bool = False,
                     use_native_client: bool = True,
                     reader_type: Optional[ReaderType] = None,
-                    buffer_size: Optional[int] = None):
+                    buffer_size: Optional[int] = None,
+                    enable_crc: bool = True):
         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, reader_type, buffer_size)
+        return cls(region, partial(gen_dataset_from_prefix, prefix), endpoint, transform, cred, client_conf,
+                   sharding, use_native_client, reader_type, buffer_size, enable_crc)
 
     def __iter__(self) -> Iterator[Any]:
         worker_id = 0

{tostorchconnector-1.0.6 → tostorchconnector-1.1.2}/tostorchconnector/tos_map_dataset.py

@@ -4,8 +4,9 @@ from typing import Any, Callable, Iterator, Optional, List, Union
 
 import torch
 
+from tosnativeclient import TosObject
 from . import TosObjectReader
-from .tos_client import CredentialProvider, TosClientConfig, TosClient, TosLogConfig, ReaderType
+from .tos_client import CredentialProvider, TosClientConfig, TosClient, ReaderType
 from .tos_common import default_trans, gen_dataset_from_prefix, \
     gen_dataset_from_urls
 from .tos_object_meta import TosObjectMeta
@@ -17,52 +18,51 @@ class TosMapDataset(torch.utils.data.Dataset):
     def __init__(self, region: str,
                  gen_dataset: Callable[[TosClient], Iterator[TosObjectMeta]],
                  endpoint: Optional[str] = None,
-                 trans: Callable[[TosObjectReader], Any] = default_trans,
+                 transform: Callable[[TosObjectReader], Any] = default_trans,
                  cred: Optional[CredentialProvider] = None,
                  client_conf: Optional[TosClientConfig] = None,
-                 log_conf: Optional[TosLogConfig] = None,
                  use_native_client: bool = True,
                  reader_type: Optional[ReaderType] = None,
-                 buffer_size: Optional[int] = None):
+                 buffer_size: Optional[int] = None,
+                 enable_crc: bool = True):
         self._gen_dataset = gen_dataset
         self._region = region
         self._endpoint = endpoint
-        self._trans = trans
+        self._trans = transform
         self._cred = cred
         self._client_conf = client_conf
-        self._log_conf = log_conf
         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)
+        self._client = TosClient(self._region, self._endpoint, self._cred, self._client_conf,
+                                 use_native_client, enable_crc)
         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,
+                  transform: Callable[[TosObjectReader], Any] = default_trans,
                   cred: Optional[CredentialProvider] = None,
                   client_conf: Optional[TosClientConfig] = None,
-                  log_conf: Optional[TosLogConfig] = None,
                   use_native_client: bool = True,
                   reader_type: Optional[ReaderType] = None,
-                  buffer_size: Optional[int] = None):
+                  buffer_size: Optional[int] = None,
+                  enable_crc: bool = True):
         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, reader_type, buffer_size)
+        return cls(region, partial(gen_dataset_from_urls, urls), endpoint, transform, cred, client_conf,
+                   use_native_client, reader_type, buffer_size, enable_crc)
 
     @classmethod
     def from_prefix(cls, prefix: str, *, region: str, endpoint: Optional[str] = None,
-                    trans: Callable[[TosObjectReader], Any] = default_trans,
+                    transform: Callable[[TosObjectReader], Any] = default_trans,
                     cred: Optional[CredentialProvider] = None,
                     client_conf: Optional[TosClientConfig] = None,
-                    log_conf: Optional[TosLogConfig] = None,
                     use_native_client: bool = True,
                     reader_type: Optional[ReaderType] = None,
-                    buffer_size: Optional[int] = None):
+                    buffer_size: Optional[int] = None,
+                    enable_crc: bool = True):
         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, reader_type, buffer_size)
+        return cls(region, partial(gen_dataset_from_prefix, prefix), endpoint, transform, cred, client_conf,
+                   use_native_client, reader_type, buffer_size, enable_crc)
 
     def __getitem__(self, i: int) -> Any:
         return self._trans_tos_object(i)

{tostorchconnector-1.0.6 → tostorchconnector-1.1.2}/tostorchconnector/tos_object_reader.py

@@ -85,7 +85,7 @@ class TosObjectStream(object):
     def close(self) -> None:
         if self._sequential_object_stream and isinstance(self._sequential_object_stream, tosnativeclient.ReadStream):
             self._sequential_object_stream.close()
-        if self._random_object_stream:
+        if self._random_object_stream and isinstance(self._random_object_stream, tosnativeclient.ReadStream):
             self._random_object_stream.close()
 
 

{tostorchconnector-1.0.6 → tostorchconnector-1.1.2/tostorchconnector.egg-info}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: tostorchconnector
-Version: 1.0.6
+Version: 1.1.2
 Summary: TOS connector integration for PyTorch
 Author-email: xiangshijian <xiangshijian@bytedance.com>
 Classifier: Development Status :: 4 - Beta
@@ -19,7 +19,7 @@ Description-Content-Type: text/markdown
 License-File: LICENSE
 Requires-Dist: torch>=2.0
 Requires-Dist: tos>=2.8.0
-Requires-Dist: tosnativeclient>=1.0.2
+Requires-Dist: tosnativeclient>=1.1.1
 Dynamic: license-file
 
 # TOS Connector for pytorch

{tostorchconnector-1.0.6 → tostorchconnector-1.1.2}/tostorchconnector.egg-info/SOURCES.txt

@@ -3,7 +3,6 @@ README.md
 pyproject.toml
 tests/test_tos_dataset.py
 tests/test_tosclient.py
-tests/test_tosrawclient.py
 tostorchconnector/__init__.py
 tostorchconnector/tos_checkpoint.py
 tostorchconnector/tos_client.py

tostorchconnector-1.1.2/tostorchconnector.egg-info/requires.txt

@@ -0,0 +1,3 @@
+torch>=2.0
+tos>=2.8.0
+tosnativeclient>=1.1.1

tostorchconnector-1.0.6/tests/test_tosrawclient.py

@@ -1,102 +0,0 @@
-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.6/tostorchconnector.egg-info/requires.txt

@@ -1,3 +0,0 @@
-torch>=2.0
-tos>=2.8.0
-tosnativeclient>=1.0.2