tostorchconnector 1.0.1__tar.gz → 1.0.3__tar.gz

{tostorchconnector-1.0.1/tostorchconnector.egg-info → tostorchconnector-1.0.3}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: tostorchconnector
-Version: 1.0.1
+Version: 1.0.3
 Summary: TOS connector integration for PyTorch
 Author-email: xiangshijian <xiangshijian@bytedance.com>
 Classifier: Development Status :: 4 - Beta

{tostorchconnector-1.0.1 → tostorchconnector-1.0.3}/pyproject.toml

@@ -5,7 +5,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "tostorchconnector"
-version = "1.0.1"
+version = "1.0.3"
 description = "TOS connector integration for PyTorch"
 authors = [{ name = "xiangshijian", email = "xiangshijian@bytedance.com" }]
 requires-python = ">=3.8,<3.13"

{tostorchconnector-1.0.1 → tostorchconnector-1.0.3}/tostorchconnector/__init__.py

@@ -3,7 +3,7 @@ from .tos_object_writer import TosObjectWriter
 from .tos_iterable_dataset import TosIterableDataset
 from .tos_map_dataset import TosMapDataset
 from .tos_checkpoint import TosCheckpoint
-from .tos_error import TosError
+from tosnativeclient import TosException, TosError
 
 __all__ = [
     'TosObjectReader',
@@ -12,4 +12,5 @@ __all__ = [
     'TosMapDataset',
     'TosCheckpoint',
     'TosError',
+    'TosException',
 ]

{tostorchconnector-1.0.1 → tostorchconnector-1.0.3}/tostorchconnector/tos_object_reader.py

@@ -2,7 +2,7 @@ import io
 import logging
 from functools import cached_property
 from os import SEEK_SET, SEEK_CUR, SEEK_END
-from typing import Optional, Callable, Iterator
+from typing import Optional, Callable, Any
 
 from tos.models2 import GetObjectOutput
 from tosnativeclient.tosnativeclient import ReadStream
@@ -16,14 +16,14 @@ class TosObjectReader(io.BufferedIOBase):
 
     def __init__(self, bucket: str, key: str,
                  get_object_meta: Optional[Callable[[], TosObjectMeta]],
-                 get_object_stream: Callable[[str, int], ReadStream | GetObjectOutput]):
+                 get_object_stream: Callable[[str, int], Any]):
         if not bucket:
             raise ValueError('bucket is empty')
         self._bucket = bucket
         self._key = key
         self._get_object_meta = get_object_meta
         self._get_object_stream = get_object_stream
-        self._object_stream: Optional[ReadStream | GetObjectOutput] = None
+        self._object_stream: Optional[Any] = None
         self._object_stream_offset = 0
         self._total_size: Optional[int] = None
         self._read_offset = 0
@@ -37,6 +37,19 @@ class TosObjectReader(io.BufferedIOBase):
     def key(self) -> str:
         return self._key
 
+    def __enter__(self):
+        return self
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        if exc_type is not None:
+            try:
+                log.info(f'Exception occurred before closing stream: {exc_type.__name__}: {exc_val}')
+            except:
+                pass
+        else:
+            if self._object_stream is not None:
+                self._object_stream.close()
+
     def read(self, size: Optional[int] = None) -> Optional[bytes]:
         if self._is_read_to_end():
             return b''
@@ -52,12 +65,15 @@ class TosObjectReader(io.BufferedIOBase):
                     while 1:
                         chunk = self._object_stream.read(self._object_stream_offset, chunk_size)
                         if not chunk:
+                            self._object_stream.close()
                             break
                         self._object_stream_offset += len(chunk)
                         self._buffer.write(chunk)
-                finally:
+                except:
                     self._object_stream.close()
-            else:
+                    raise
+
+            elif isinstance(self._object_stream, GetObjectOutput):
                 for chunk in self._object_stream:
                     self._buffer.write(chunk)
 
@@ -149,13 +165,15 @@ class TosObjectReader(io.BufferedIOBase):
                 while offset > size:
                     chunk = self._object_stream.read(self._object_stream_offset, chunk_size)
                     if not chunk:
+                        self._object_stream.close()
                         break
                     size += self._buffer.write(chunk)
                     self._object_stream_offset += len(chunk)
                 self._total_size = self._buffer.tell()
-            finally:
+            except:
                 self._object_stream.close()
-        else:
+                raise
+        elif isinstance(self._object_stream, GetObjectOutput):
             try:
                 while offset > size:
                     size += self._buffer.write(next(self._object_stream))

{tostorchconnector-1.0.1 → tostorchconnector-1.0.3}/tostorchconnector/tos_object_writer.py

@@ -1,10 +1,9 @@
 import io
 import logging
 import threading
-from typing import Callable, Optional, Any
+from typing import Callable, Any
 
 from tos.models2 import PutObjectOutput
-from tosnativeclient.tosnativeclient import WriteStream
 
 log = logging.getLogger(__name__)
 
@@ -24,7 +23,7 @@ class PutObjectStream(object):
 
 class TosObjectWriter(io.BufferedIOBase):
 
-    def __init__(self, bucket: str, key: str, put_object_stream: WriteStream | PutObjectStream):
+    def __init__(self, bucket: str, key: str, put_object_stream: Any):
         if not bucket:
             raise ValueError('bucket is empty')
         self._bucket = bucket

{tostorchconnector-1.0.1 → tostorchconnector-1.0.3/tostorchconnector.egg-info}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: tostorchconnector
-Version: 1.0.1
+Version: 1.0.3
 Summary: TOS connector integration for PyTorch
 Author-email: xiangshijian <xiangshijian@bytedance.com>
 Classifier: Development Status :: 4 - Beta

{tostorchconnector-1.0.1 → tostorchconnector-1.0.3}/tostorchconnector.egg-info/SOURCES.txt

@@ -7,7 +7,6 @@ tostorchconnector/__init__.py
 tostorchconnector/tos_checkpoint.py
 tostorchconnector/tos_client.py
 tostorchconnector/tos_common.py
-tostorchconnector/tos_error.py
 tostorchconnector/tos_iterable_dataset.py
 tostorchconnector/tos_map_dataset.py
 tostorchconnector/tos_object_meta.py

tostorchconnector-1.0.1/tostorchconnector/tos_error.py

@@ -1,11 +0,0 @@
-from typing import Optional
-
-
-class TosError(Exception):
-    def __init__(self, message: str, status_code: Optional[int] = None, ec: Optional[str] = None,
-                 request_id: Optional[str] = None):
-        super().__init__(message)
-        self.message = message
-        self.status_code = status_code
-        self.ec = ec
-        self.request_id = request_id