tostorchconnector 1.0.3__tar.gz → 1.0.4__tar.gz

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

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

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

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

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

@@ -1,5 +1,6 @@
 import io
 import logging
+import threading
 from functools import cached_property
 from os import SEEK_SET, SEEK_CUR, SEEK_END
 from typing import Optional, Callable, Any
@@ -28,6 +29,8 @@ class TosObjectReader(io.BufferedIOBase):
         self._total_size: Optional[int] = None
         self._read_offset = 0
         self._buffer = io.BytesIO()
+        self._closed = False
+        self._lock = threading.Lock()
 
     @property
     def bucket(self) -> str:
@@ -47,7 +50,14 @@ class TosObjectReader(io.BufferedIOBase):
             except:
                 pass
         else:
-            if self._object_stream is not None:
+            self.close()
+
+    def close(self) -> None:
+        with self._lock:
+            if self._closed:
+                return
+            self._closed = True
+            if self._object_stream and isinstance(self._object_stream, ReadStream):
                 self._object_stream.close()
 
     def read(self, size: Optional[int] = None) -> Optional[bytes]:
@@ -65,12 +75,12 @@ class TosObjectReader(io.BufferedIOBase):
                     while 1:
                         chunk = self._object_stream.read(self._object_stream_offset, chunk_size)
                         if not chunk:
-                            self._object_stream.close()
+                            self.close()
                             break
                         self._object_stream_offset += len(chunk)
                         self._buffer.write(chunk)
                 except:
-                    self._object_stream.close()
+                    self.close()
                     raise
 
             elif isinstance(self._object_stream, GetObjectOutput):
@@ -165,13 +175,13 @@ 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()
+                        self.close()
                         break
                     size += self._buffer.write(chunk)
                     self._object_stream_offset += len(chunk)
                 self._total_size = self._buffer.tell()
             except:
-                self._object_stream.close()
+                self.close()
                 raise
         elif isinstance(self._object_stream, GetObjectOutput):
             try:

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

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