tostorchconnector 1.0.7__tar.gz → 1.0.9__tar.gz

{tostorchconnector-1.0.7/tostorchconnector.egg-info → tostorchconnector-1.0.9}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: tostorchconnector
-Version: 1.0.7
+Version: 1.0.9
 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.4
+Requires-Dist: tosnativeclient>=1.0.6
 Dynamic: license-file
 
 # TOS Connector for pytorch

{tostorchconnector-1.0.7 → tostorchconnector-1.0.9}/pyproject.toml

@@ -5,7 +5,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "tostorchconnector"
-version = "1.0.7"
+version = "1.0.9"
 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.4"
+    "tosnativeclient >= 1.0.6"
 ]
 
 [tool.setuptools.packages]

{tostorchconnector-1.0.7 → tostorchconnector-1.0.9}/tests/test_tos_dataset.py

@@ -1,4 +1,5 @@
 import os
+import pickle
 import unittest
 
 from torch.utils.data import DataLoader
@@ -7,6 +8,7 @@ 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):
@@ -24,24 +26,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')
@@ -49,17 +77,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')
@@ -75,6 +109,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.7 → tostorchconnector-1.0.9}/tests/test_tosclient.py

@@ -51,9 +51,9 @@ 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:

{tostorchconnector-1.0.7 → tostorchconnector-1.0.9}/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'
@@ -79,40 +119,62 @@ 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
+        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._log_conf = TosLogConfig() if log_conf is None else log_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
+
+    @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:
+                    directives = ''
+                    directory = ''
+                    file_name_prefix = ''
+                    if self._log_conf.log_dir and self._log_conf.log_file_name:
+                        if self._log_conf.log_level:
+                            if self._log_conf.log_level == logging.DEBUG:
+                                directives = 'debug'
+                            elif self._log_conf.log_level == logging.INFO:
+                                directives = 'info'
+                            elif self._log_conf.log_level == logging.WARN:
+                                directives = 'warn'
+                            elif self._log_conf.log_level == logging.ERROR:
+                                directives = 'error'
+                            else:
+                                directives = 'info'
+                        directory = self._log_conf.log_dir
+                        file_name_prefix = self._log_conf.log_file_name
+                        # reset log_conf to avoid panic
+                        self._log_conf = TosLogConfig()
+
+                    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,
+                                                                   directives=directives,
+                                                                   directory=directory,
+                                                                   file_name_prefix=file_name_prefix)
+                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)
+                    if self._log_conf.log_dir and self._log_conf.log_file_name:
+                        file_path = os.path.join(self._log_conf.log_dir, self._log_conf.log_file_name)
+                        log_level = self._log_conf.log_level if self._log_conf.log_level else logging.INFO
+                        tos.set_logger(file_path=file_path, level=log_level)
+
+                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 +217,15 @@ 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) -> 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)
         raise NotImplementedError()
 
     def list_objects(self, bucket: str, prefix: str, max_keys: int = 1000,

{tostorchconnector-1.0.7 → tostorchconnector-1.0.9}/tostorchconnector/tos_common.py

@@ -8,19 +8,28 @@ from .tos_object_meta import TosObjectMeta
 log = logging.getLogger(__name__)
 
 
+class TosObjectIterable(object):
+    def __init__(self, bucket: str, prefix: str, client: TosClient):
+        self._bucket = bucket
+        self._prefix = prefix
+        self._client = client
+
+    def __iter__(self) -> Iterator[TosObjectMeta]:
+        return iter(TosObjectIterator(self._bucket, self._prefix, self._client))
+
+
 class TosObjectIterator(object):
     def __init__(self, bucket: str, prefix: str, client: TosClient):
         self._bucket = bucket
         self._prefix = prefix
         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,13 +42,15 @@ 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)
 
             if self._object_metas is None or self._index >= len(self._object_metas):
                 self._object_metas = []
                 self._index = 0
                 while 1:
                     objects = next(self._list_stream)
+                    self._continuation_token = self._list_stream.current_continuation_token()
                     for content in objects.contents:
                         self._object_metas.append(
                             TosObjectMeta(content.bucket, content.key, content.size, content.etag))
@@ -101,4 +112,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.7 → tostorchconnector-1.0.9}/tostorchconnector/tos_iterable_dataset.py

@@ -16,7 +16,7 @@ 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,
@@ -27,7 +27,7 @@ class TosIterableDataset(torch.utils.data.IterableDataset):
         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

{tostorchconnector-1.0.7 → tostorchconnector-1.0.9}/tostorchconnector/tos_map_dataset.py

@@ -17,7 +17,7 @@ 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,
@@ -27,7 +27,7 @@ class TosMapDataset(torch.utils.data.Dataset):
         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
@@ -40,7 +40,7 @@ class TosMapDataset(torch.utils.data.Dataset):
 
     @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,
@@ -48,12 +48,12 @@ class TosMapDataset(torch.utils.data.Dataset):
                   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,
+        return cls(region, partial(gen_dataset_from_urls, urls), endpoint, transform, cred, client_conf, log_conf,
                    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,
+                    transform: Callable[[TosObjectReader], Any] = default_trans,
                     cred: Optional[CredentialProvider] = None,
                     client_conf: Optional[TosClientConfig] = None,
                     log_conf: Optional[TosLogConfig] = None,
@@ -61,7 +61,7 @@ class TosMapDataset(torch.utils.data.Dataset):
                     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,
+        return cls(region, partial(gen_dataset_from_prefix, prefix), endpoint, transform, cred, client_conf, log_conf,
                    use_native_client, reader_type, buffer_size)
 
     def __getitem__(self, i: int) -> Any:

{tostorchconnector-1.0.7 → tostorchconnector-1.0.9}/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.7 → tostorchconnector-1.0.9/tostorchconnector.egg-info}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: tostorchconnector
-Version: 1.0.7
+Version: 1.0.9
 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.4
+Requires-Dist: tosnativeclient>=1.0.6
 Dynamic: license-file
 
 # TOS Connector for pytorch

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

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

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

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