tostorchconnector 1.0.4__py3-none-any.whl → 1.0.5__py3-none-any.whl

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/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/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/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/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/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)

tostorchconnector/tos_object_reader.py

@@ -1,34 +1,103 @@
 import io
 import logging
 import threading
-from functools import cached_property
+from abc import ABC, abstractmethod
+from functools import partial, cached_property
 from os import SEEK_SET, SEEK_CUR, SEEK_END
 from typing import Optional, Callable, Any
 
-from tos.models2 import GetObjectOutput
-from tosnativeclient.tosnativeclient import ReadStream
-
+import tosnativeclient
 from .tos_object_meta import TosObjectMeta
 
 log = logging.getLogger(__name__)
 
+DEFAULT_CHUNK_SIZE = 1 * 1024 * 1024
+DEFAULT_BUFFER_SIZE = 8 * 1024 * 1024
 
-class TosObjectReader(io.BufferedIOBase):
 
-    def __init__(self, bucket: str, key: str,
-                 get_object_meta: Optional[Callable[[], TosObjectMeta]],
-                 get_object_stream: Callable[[str, int], Any]):
+class TosObjectStream(object):
+    def __init__(self, bucket: str, key: str, get_object_meta: Optional[Callable[[], TosObjectMeta]],
+                 client: Any):
+        self._bucket = bucket
+        self._key = key
+        self._get_object_meta = get_object_meta
+        self._client = client
+        self._sequential_object_stream = None
+        self._sequential_object_stream_offset = 0
+        self._random_object_stream = None
+
+    def sequential_read(self, chunk_size) -> Optional[bytes]:
+        self._trigger_prefetch()
+        assert self._sequential_object_stream is not None
+        if chunk_size <= 0:
+            chunk_size = DEFAULT_CHUNK_SIZE
+        if isinstance(self._sequential_object_stream, tosnativeclient.ReadStream):
+            chunk = self._sequential_object_stream.read(self._sequential_object_stream_offset, chunk_size)
+            if chunk:
+                self._sequential_object_stream_offset += len(chunk)
+        else:
+            chunk = self._sequential_object_stream.read(chunk_size)
+        return chunk
+
+    def random_read(self, read_start, read_end, chunk_size, callback: Callable[[bytes], None]) -> None:
+        if chunk_size <= 0:
+            chunk_size = DEFAULT_CHUNK_SIZE
+        if isinstance(self._client, tosnativeclient.TosClient):
+            if self._random_object_stream is None:
+                object_meta = self.object_meta
+                self._random_object_stream = self._client.get_object(self._bucket, self._key, object_meta.etag,
+                                                                     object_meta.size)
+
+            offset = read_start
+            while 1:
+                if offset >= read_end:
+                    break
+                length = chunk_size if offset + chunk_size <= read_end else read_end - offset
+                chunk = self._random_object_stream.read(offset, length)
+                if not chunk:
+                    break
+                callback(chunk)
+                offset += len(chunk)
+            return
+
+        object_meta = self.object_meta
+        output = self._client.get_object(self._bucket, self._key, if_match=object_meta.etag,
+                                         range=f'bytes={read_start}-{read_end - read_start + 1}')
+        while 1:
+            chunk = output.read(chunk_size)
+            if not chunk:
+                break
+            callback(chunk)
+
+    def _trigger_prefetch(self) -> None:
+        if self._sequential_object_stream is None:
+            object_meta = self.object_meta
+            if isinstance(self._client, tosnativeclient.TosClient):
+                get_object_stream = partial(self._client.get_object, self._bucket, self._key)
+            else:
+                get_object_stream = lambda et, sz: self._client.get_object(self._bucket, self._key, '', et)
+            self._sequential_object_stream = get_object_stream(object_meta.etag, object_meta.size)
+
+    @cached_property
+    def object_meta(self) -> TosObjectMeta:
+        return self._get_object_meta()
+
+    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:
+            self._random_object_stream.close()
+
+
+class TosObjectReader(ABC, io.BufferedIOBase):
+    def __init__(self, bucket: str, key: str, object_stream: TosObjectStream):
         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[Any] = None
-        self._object_stream_offset = 0
+        self._object_stream = object_stream
         self._total_size: Optional[int] = None
         self._read_offset = 0
-        self._buffer = io.BytesIO()
         self._closed = False
         self._lock = threading.Lock()
 
@@ -40,10 +109,22 @@ class TosObjectReader(io.BufferedIOBase):
     def key(self) -> str:
         return self._key
 
+    @property
+    def closed(self) -> bool:
+        return self._closed
+
+    def close(self) -> None:
+        if self._closed:
+            return
+        with self._lock:
+            if not self._closed:
+                self._closed = True
+                self._object_stream.close()
+
     def __enter__(self):
         return self
 
-    def __exit__(self, exc_type, exc_val, exc_tb):
+    def __exit__(self, exc_type, exc_val, exc_tb) -> None:
         if exc_type is not None:
             try:
                 log.info(f'Exception occurred before closing stream: {exc_type.__name__}: {exc_val}')
@@ -52,41 +133,62 @@ class TosObjectReader(io.BufferedIOBase):
         else:
             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()
+    @abstractmethod
+    def read(self, size: Optional[int] = None) -> bytes:
+        pass
+
+    @abstractmethod
+    def readinto(self, buf) -> int:
+        pass
+
+    @abstractmethod
+    def seek(self, offset: int, whence: int = SEEK_SET) -> int:
+        pass
+
+    def tell(self) -> int:
+        return self._read_offset
+
+    def readable(self) -> bool:
+        return self.closed
+
+    def writable(self) -> bool:
+        return False
+
+    def seekable(self) -> bool:
+        return True
+
+    def _is_read_to_end(self) -> bool:
+        if self._total_size is None:
+            return False
+        return self._read_offset == self._total_size
+
+    def _get_total_size(self) -> int:
+        if self._total_size is None:
+            self._total_size = self._object_stream.object_meta.size
+        return self._total_size
+
+
+class SequentialTosObjectReader(TosObjectReader):
 
-    def read(self, size: Optional[int] = None) -> Optional[bytes]:
+    def __init__(self, bucket: str, key: str, object_stream: TosObjectStream):
+        super().__init__(bucket, key, object_stream)
+        self._buffer = io.BytesIO()
+
+    def read(self, size: Optional[int] = None) -> bytes:
         if self._is_read_to_end():
             return b''
 
-        self._trigger_prefetch()
         current_read_offset = self._read_offset
         if size is None or size < 0:
+            if self.closed:
+                raise RuntimeError('read on closed GetObjectOutput')
             # means read all
             self._buffer.seek(0, SEEK_END)
-            if isinstance(self._object_stream, ReadStream):
-                try:
-                    chunk_size = 1 * 1024 * 1024
-                    while 1:
-                        chunk = self._object_stream.read(self._object_stream_offset, chunk_size)
-                        if not chunk:
-                            self.close()
-                            break
-                        self._object_stream_offset += len(chunk)
-                        self._buffer.write(chunk)
-                except:
-                    self.close()
-                    raise
-
-            elif isinstance(self._object_stream, GetObjectOutput):
-                for chunk in self._object_stream:
-                    self._buffer.write(chunk)
-
+            while 1:
+                chunk = self._object_stream.sequential_read(DEFAULT_CHUNK_SIZE)
+                if not chunk:
+                    break
+                self._buffer.write(chunk)
             self._total_size = self._buffer.tell()
         else:
             self.seek(size, SEEK_CUR)
@@ -96,12 +198,10 @@ class TosObjectReader(io.BufferedIOBase):
         self._read_offset = self._buffer.tell()
         return data
 
-    def readinto(self, buf) -> Optional[int]:
+    def readinto(self, buf) -> int:
         size = len(buf)
         if self._is_read_to_end() or size == 0:
             return 0
-
-        self._trigger_prefetch()
         current_read_offset = self._read_offset
         self.seek(size, SEEK_CUR)
         self._buffer.seek(current_read_offset)
@@ -123,7 +223,7 @@ class TosObjectReader(io.BufferedIOBase):
         elif whence == SEEK_SET:
             pass
         else:
-            raise ValueError('invalid whence')
+            raise ValueError('invalid whence, must be passed SEEK_CUR, SEEK_SET, or SEEK_END')
 
         if offset < 0:
             raise ValueError(f'invalid seek offset {offset}')
@@ -131,68 +231,156 @@ class TosObjectReader(io.BufferedIOBase):
         if offset > self._buffer_size():
             self._prefetch_to_offset(offset)
 
+        if self._total_size is not None:
+            offset = min(offset, self._total_size)
+
         self._read_offset = self._buffer.seek(offset)
         return self._read_offset
 
-    def tell(self) -> int:
+    def _prefetch_to_offset(self, offset: int) -> None:
+        if self.closed:
+            raise RuntimeError('read on closed GetObjectOutput')
+        size = self._buffer.seek(0, SEEK_END)
+        while offset > size:
+            chunk = self._object_stream.sequential_read(DEFAULT_CHUNK_SIZE)
+            if not chunk:
+                break
+            size += self._buffer.write(chunk)
+
+        self._total_size = self._buffer.tell()
+
+    def _buffer_size(self) -> int:
+        cur_pos = self._buffer.tell()
+        self._buffer.seek(0, SEEK_END)
+        buffer_size = self._buffer.tell()
+        self._buffer.seek(cur_pos)
+        return buffer_size
+
+
+class RangedTosObjectReader(TosObjectReader):
+    def __init__(self, bucket: str, key: str,
+                 object_stream: TosObjectStream, buffer_size: Optional[int] = None):
+        super().__init__(bucket, key, object_stream)
+
+        if buffer_size is None:
+            self._buffer_size = DEFAULT_BUFFER_SIZE
+            self._enable_buffering = True
+        else:
+            self._buffer_size = buffer_size
+            self._enable_buffering = buffer_size > 0
+
+        self._buffer = bytearray(self._buffer_size) if self._enable_buffering else None
+        self._buffer_view = memoryview(self._buffer) if self._buffer else None
+        self._buffer_start = 0
+        self._buffer_end = 0
+
+    def read(self, size: Optional[int] = None) -> bytes:
+        if self._is_read_to_end():
+            return b''
+
+        read_start = self._read_offset
+        if size is None or size < 0:
+            read_end = self._get_total_size()
+        else:
+            read_end = min(read_start + size, self._get_total_size())
+
+        if read_start >= read_end:
+            return b''
+
+        view = memoryview(bytearray(read_end - read_start))
+        self._read_into_view(view, read_start, read_end)
+        return view.tobytes()
+
+    def readinto(self, buf) -> int:
+        size = len(buf)
+        if self._is_read_to_end() or size == 0:
+            return 0
+
+        try:
+            view = memoryview(buf)
+            if view.readonly:
+                raise TypeError(f'argument must be a writable bytes-like object, not {type(buf).__name__}')
+        except TypeError:
+            raise TypeError(f'argument must be a writable bytes-like object, not {type(buf).__name__}')
+
+        read_start = self._read_offset
+        read_end = min(read_start + size, self._get_total_size())
+        if read_start >= read_end:
+            return 0
+
+        return self._read_into_view(view, read_start, read_end)
+
+    def seek(self, offset: int, whence: int = SEEK_SET) -> int:
+        if whence == SEEK_END:
+            if offset >= 0:
+                self._read_offset = self._get_total_size()
+                return self._read_offset
+            # offset is negative
+            offset += self._get_total_size()
+        elif whence == SEEK_CUR:
+            if self._is_read_to_end() and offset >= 0:
+                return self._read_offset
+            offset += self._read_offset
+        elif whence == SEEK_SET:
+            pass
+        else:
+            raise ValueError('invalid whence, must be passed SEEK_CUR, SEEK_SET, or SEEK_END')
+
+        if offset < 0:
+            raise ValueError(f'invalid seek offset {offset}')
+
+        self._read_offset = min(offset, self._get_total_size())
         return self._read_offset
 
-    def readable(self):
-        return True
+    def _read_into_view(self, view: memoryview, read_start: int, read_end: int) -> int:
+        readed = 0
+        if self._buffer_start <= read_start < self._buffer_end <= read_end:
+            readed_once = self._read_from_buffer(view, read_start, self._buffer_end)
+            read_start = self._buffer_end
+            view = view[readed_once:]
+            readed += readed_once
 
-    def writable(self):
-        return False
+        if read_end - read_start >= self._buffer_size or not self._enable_buffering:
+            readed += self._read_directly(view, read_start, read_end)
+        else:
+            readed += self._read_from_buffer(view, read_start, read_end)
 
-    def seekable(self):
-        return True
+        self._read_offset += readed
+        return readed
 
-    @cached_property
-    def _object_meta(self) -> TosObjectMeta:
-        return self._get_object_meta()
+    def _read_directly(self, view: memoryview, read_start: int, read_end: int) -> int:
+        readed = 0
 
-    def _trigger_prefetch(self) -> None:
-        if self._object_stream is None:
-            object_meta = self._object_meta
-            self._object_stream = self._get_object_stream(object_meta.etag, object_meta.size)
-            self._object_stream_offset = 0
+        def callback(data: bytes):
+            nonlocal readed
+            view[readed: readed + len(data)] = data
+            readed += len(data)
 
-    def _is_read_to_end(self) -> bool:
-        if self._total_size is None:
-            return False
-        return self._read_offset == self._total_size
+        self._object_stream.random_read(read_start, read_end, DEFAULT_CHUNK_SIZE, callback)
+        return readed
 
-    def _get_total_size(self) -> int:
-        if self._total_size is None:
-            self._total_size = self._object_meta.size
-        return self._total_size
+    def _read_from_buffer(self, view: memoryview, read_start: int, read_end: int) -> int:
+        if read_start < self._buffer_start or read_end > self._buffer_end:
+            self._load_buffer(read_start)
 
-    def _prefetch_to_offset(self, offset: int) -> None:
-        self._trigger_prefetch()
-        size = self._buffer.seek(0, SEEK_END)
-        if isinstance(self._object_stream, ReadStream):
-            try:
-                chunk_size = 1 * 1024 * 1024
-                while offset > size:
-                    chunk = self._object_stream.read(self._object_stream_offset, chunk_size)
-                    if not chunk:
-                        self.close()
-                        break
-                    size += self._buffer.write(chunk)
-                    self._object_stream_offset += len(chunk)
-                self._total_size = self._buffer.tell()
-            except:
-                self.close()
-                raise
-        elif isinstance(self._object_stream, GetObjectOutput):
-            try:
-                while offset > size:
-                    size += self._buffer.write(next(self._object_stream))
-            except StopIteration:
-                self._total_size = self._buffer.tell()
+        buffer_offset = read_start - self._buffer_start
+        readed = read_end - read_start
+        assert self._buffer is not None
+        view[:readed] = self._buffer[buffer_offset:buffer_offset + readed]
+        return readed
 
-    def _buffer_size(self) -> int:
-        cur_pos = self._buffer.tell()
-        self._buffer.seek(0, SEEK_END)
-        buffer_size = self._buffer.tell()
-        self._buffer.seek(cur_pos)
-        return buffer_size
+    def _load_buffer(self, read_start: int) -> None:
+        read_end = min(read_start + self._buffer_size, self._get_total_size())
+        assert self._buffer_view is not None
+
+        readed = 0
+
+        def callback(data: bytes):
+            nonlocal readed
+            self._buffer_view[readed: readed + len(data)] = data
+            readed += len(data)
+
+        self._object_stream.random_read(read_start, read_end, DEFAULT_CHUNK_SIZE, callback)
+
+        self._buffer_start = read_start
+        self._buffer_end = read_start + readed

tostorchconnector/tos_object_writer.py

@@ -1,7 +1,7 @@
 import io
 import logging
 import threading
-from typing import Callable, Any
+from typing import Callable, Any, Union
 
 from tos.models2 import PutObjectOutput
 
@@ -12,11 +12,16 @@ class PutObjectStream(object):
     def __init__(self, put_object: Callable[[io.BytesIO], PutObjectOutput]):
         self._put_object = put_object
         self._buffer = io.BytesIO()
+        self._closed = False
 
-    def write(self, data):
+    def write(self, data) -> int:
+        if self._closed:
+            raise RuntimeError('write on closed PutObjectStream')
         self._buffer.write(data)
+        return len(data)
 
     def close(self):
+        self._closed = True
         self._buffer.seek(0)
         _ = self._put_object(self._buffer)
 
@@ -41,11 +46,15 @@ class TosObjectWriter(io.BufferedIOBase):
     def key(self) -> str:
         return self._key
 
+    @property
+    def closed(self) -> bool:
+        return self._closed
+
     def __enter__(self):
         self._write_offset = 0
         return self
 
-    def __exit__(self, exc_type, exc_val, exc_tb):
+    def __exit__(self, exc_type, exc_val, exc_tb) -> None:
         if exc_type is not None:
             try:
                 log.info(f'Exception occurred before closing stream: {exc_type.__name__}: {exc_val}')
@@ -54,17 +63,22 @@ class TosObjectWriter(io.BufferedIOBase):
         else:
             self.close()
 
-    def write(self, data) -> int:
-        self._put_object_stream.write(data)
-        self._write_offset += len(data)
-        return len(data)
+    def write(self, data: Union[bytes, memoryview]) -> int:
+        if isinstance(data, memoryview):
+            data = data.tobytes()
+        written = self._put_object_stream.write(data)
+        assert written == len(data)
+        self._write_offset += written
+        return written
 
     def close(self) -> None:
+        if self._closed:
+            return
+
         with self._lock:
-            if self._closed:
-                return
-            self._closed = True
-            self._put_object_stream.close()
+            if not self._closed:
+                self._closed = True
+                self._put_object_stream.close()
 
     def tell(self) -> int:
         return self._write_offset
@@ -72,11 +86,11 @@ class TosObjectWriter(io.BufferedIOBase):
     def flush(self) -> None:
         pass
 
-    def readable(self):
+    def readable(self) -> bool:
         return False
 
-    def writable(self):
-        return True
+    def writable(self) -> bool:
+        return not self.closed
 
-    def seekable(self):
+    def seekable(self) -> bool:
         return False

{tostorchconnector-1.0.4.dist-info → tostorchconnector-1.0.5.dist-info}/METADATA

@@ -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.5.dist-info/RECORD

@@ -0,0 +1,14 @@
+tostorchconnector/__init__.py,sha256=GASYEQEKI6nWayop8F8Ln-k7NY4jpmmV9qSy_-gY5Tk,508
+tostorchconnector/tos_checkpoint.py,sha256=xbBWHeESuI6mml5VPPFwUwkE5EZYSRGYeeXTNYhaRQY,1400
+tostorchconnector/tos_client.py,sha256=a8JkE41tsa-77RaCPWjq6mIY0n2AuW6fWkY2Ya1T4r8,7364
+tostorchconnector/tos_common.py,sha256=1TRHJdCQjlc_mSOI5j2VoxXGH-Xbamc9bCd0SZyB0L8,3372
+tostorchconnector/tos_iterable_dataset.py,sha256=EMLNzsH4E1qUFcN6S9bhGBbA4bRAAcGY9o5Z7sHYFIk,4796
+tostorchconnector/tos_map_dataset.py,sha256=4fSLR8iPA8jXcMD94teEQv2h6K3MG35QGC119_SWRds,3913
+tostorchconnector/tos_object_meta.py,sha256=YrEQikioD5v0C_VcoudvTn0apUcxqxmSNQ4dDIip1Zc,562
+tostorchconnector/tos_object_reader.py,sha256=a0KCw70hokWi00MEssjvaGaL-WxM040_lxtNQ8l9XEU,13517
+tostorchconnector/tos_object_writer.py,sha256=GfQUHZ3RLmaJzqOtov4EwgTl6eI3IafhJ2xnFK9cZ-k,2429
+tostorchconnector-1.0.5.dist-info/licenses/LICENSE,sha256=WBJyZyF8q8ZxohuLttIbv7HNLBRA0NnnRWWgfDlDZBE,11361
+tostorchconnector-1.0.5.dist-info/METADATA,sha256=vhpr7HXeZ4vPT9yNnL1uk6KGIYephGJsdgf6CxHlUCA,932
+tostorchconnector-1.0.5.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+tostorchconnector-1.0.5.dist-info/top_level.txt,sha256=w9WvBP6KEi5Dourf86IdBbpWe-fad84S4SBftSg7H-k,18
+tostorchconnector-1.0.5.dist-info/RECORD,,

tostorchconnector-1.0.4.dist-info/RECORD

@@ -1,14 +0,0 @@
-tostorchconnector/__init__.py,sha256=RX4fIC9N-GR5THq9bnfzPPf4XIpTQy4OGfdhpmZz2UU,448
-tostorchconnector/tos_checkpoint.py,sha256=SHrPTZ_HPKG3OvIthj2AXPwyKQTTUxEBYnaNNEZrSj0,1497
-tostorchconnector/tos_client.py,sha256=8artQiMoObamN7gGYVdjsSfoPTwmBVsLTt5Y898ohCc,7165
-tostorchconnector/tos_common.py,sha256=oxxxzpA7Adn3hwNkpB9fHCueGGmZEc2f9tFqtEgPn9M,3167
-tostorchconnector/tos_iterable_dataset.py,sha256=a28XPfZJwvWJJeI3qhHXJAAhKyg4U1DHavUI3VsQorY,4312
-tostorchconnector/tos_map_dataset.py,sha256=o8Fq8O5aG2LS8NA9eFtEfBbkisWwHg-dItO8UAqu08o,3471
-tostorchconnector/tos_object_meta.py,sha256=YrEQikioD5v0C_VcoudvTn0apUcxqxmSNQ4dDIip1Zc,562
-tostorchconnector/tos_object_reader.py,sha256=g2Iek4eDQllbr5AeFe81kB_zKSBBw18BGYElZqbnPcE,6472
-tostorchconnector/tos_object_writer.py,sha256=-I2coKCuUITL_LHf5we3cZBQ6rdn0bQL6NwpF39uDuU,1949
-tostorchconnector-1.0.4.dist-info/licenses/LICENSE,sha256=WBJyZyF8q8ZxohuLttIbv7HNLBRA0NnnRWWgfDlDZBE,11361
-tostorchconnector-1.0.4.dist-info/METADATA,sha256=wQTa2P1xE5nSaoyHGABTHHRF1lqx29mQeD6OyoXWnE0,881
-tostorchconnector-1.0.4.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
-tostorchconnector-1.0.4.dist-info/top_level.txt,sha256=w9WvBP6KEi5Dourf86IdBbpWe-fad84S4SBftSg7H-k,18
-tostorchconnector-1.0.4.dist-info/RECORD,,