tostorchconnector 1.0.0__tar.gz → 1.0.1__tar.gz

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

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

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

@@ -5,7 +5,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "tostorchconnector"
-version = "1.0.0"
+version = "1.0.1"
 description = "TOS connector integration for PyTorch"
 authors = [{ name = "xiangshijian", email = "xiangshijian@bytedance.com" }]
 requires-python = ">=3.8,<3.13"
@@ -29,4 +29,4 @@ dependencies = [
 ]
 
 [tool.setuptools.packages]
-find = { where = ["tostorchconnector"] }
+find = { where = ["."], include = ["tostorchconnector"] }

tostorchconnector-1.0.1/tostorchconnector/__init__.py

@@ -0,0 +1,15 @@
+from .tos_object_reader import TosObjectReader
+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
+
+__all__ = [
+    'TosObjectReader',
+    'TosObjectWriter',
+    'TosIterableDataset',
+    'TosMapDataset',
+    'TosCheckpoint',
+    'TosError',
+]

tostorchconnector-1.0.1/tostorchconnector/tos_checkpoint.py

@@ -0,0 +1,39 @@
+import logging
+from typing import Optional
+
+from . import TosObjectReader, TosObjectWriter
+from .tos_client import CredentialProvider, TosClientConfig, TosClient, TosLogConfig
+from .tos_common import parse_tos_url
+
+log = logging.getLogger(__name__)
+
+
+class TosCheckpoint(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=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
+
+    def reader(self, url: str) -> TosObjectReader:
+        bucket, key = parse_tos_url(url)
+        return self._get_tos_client().get_object(bucket, key)
+
+    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

tostorchconnector-1.0.1/tostorchconnector/tos_client.py

@@ -0,0 +1,169 @@
+import io
+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 .tos_object_meta import TosObjectMeta
+from .tos_object_writer import PutObjectStream
+
+log = logging.getLogger(__name__)
+
+
+class CredentialProvider(object):
+    def __init__(self, ak: str, sk: str):
+        self._ak = ak
+        self._sk = sk
+
+    @property
+    def ak(self) -> str:
+        return self._ak
+
+    @property
+    def sk(self) -> str:
+        return self._sk
+
+
+class TosClientConfig(object):
+    def __init__(self, part_size: int = 8 * 1024 * 1024,
+                 max_retry_count: int = 3, shared_prefetch_tasks: int = 20):
+        self._part_size = part_size
+        self._max_retry_count = max_retry_count
+        self._shared_prefetch_tasks = shared_prefetch_tasks
+
+    @property
+    def part_size(self) -> int:
+        return self._part_size
+
+    @property
+    def max_retry_count(self) -> int:
+        return self._max_retry_count
+
+    @property
+    def shared_prefetch_tasks(self) -> int:
+        return self._shared_prefetch_tasks
+
+
+class TosLogConfig(object):
+    def __init__(self, log_dir: str = '',
+                 log_file_name: str = '', log_level: Optional[int] = logging.INFO):
+        self._log_dir = log_dir
+        self._log_file_name = log_file_name
+        self._log_level = log_level
+
+    @property
+    def log_level(self) -> Optional[int]:
+        return self._log_level
+
+    @property
+    def log_dir(self) -> str:
+        return self._log_dir
+
+    @property
+    def log_file_name(self) -> str:
+        return self._log_file_name
+
+
+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._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)
+
+    @property
+    def use_native_client(self):
+        return self._use_native_client
+
+    def get_object(self, bucket: str, key: str, etag: Optional[str] = None,
+                   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)
+
+    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)
+        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)
+
+        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:
+        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)
+        raise NotImplementedError()
+
+    def list_objects(self, bucket: str, prefix: str, max_keys: int = 1000,
+                     continuation_token: Optional[str] = None, delimiter: Optional[str] = None) -> Tuple[
+        List[TosObjectMeta], bool, Optional[str]]:
+        log.debug(f'list_objects tos://{bucket}/{prefix}')
+        resp = self._client.list_objects_type2(bucket, prefix, max_keys=max_keys, continuation_token=continuation_token,
+                                               delimiter=delimiter)
+        object_metas = []
+        for obj in resp.contents:
+            object_metas.append(TosObjectMeta(bucket, obj.key, obj.size, obj.etag))
+        return object_metas, resp.is_truncated, resp.next_continuation_token

tostorchconnector-1.0.1/tostorchconnector/tos_common.py

@@ -0,0 +1,99 @@
+import logging
+from typing import Union, Iterator, Tuple, Optional
+
+from . import TosObjectReader
+from .tos_client import TosClient
+from .tos_object_meta import TosObjectMeta
+
+log = logging.getLogger(__name__)
+
+
+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._object_metas = None
+        self._index = 0
+
+        self._is_truncated = True
+        self._continuation_token = None
+
+    def close(self):
+        if self._list_stream is not None:
+            self._list_stream.close()
+
+    def __iter__(self) -> Iterator[TosObjectMeta]:
+        return self
+
+    def __next__(self):
+        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)
+
+            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))
+
+            object_meta = self._object_metas[self._index]
+            self._index += 1
+            return object_meta
+
+        while self._object_metas is None or self._index >= len(self._object_metas):
+            if not self._is_truncated:
+                raise StopIteration
+            self._object_metas, self._is_truncated, self._continuation_token = self._client.list_objects(
+                self._bucket,
+                self._prefix,
+                max_keys=1000,
+                continuation_token=self._continuation_token,
+                delimiter=self._delimiter)
+            self._index = 0
+
+        object_meta = self._object_metas[self._index]
+        self._index += 1
+        return object_meta
+
+
+def parse_tos_url(url: str) -> Tuple[str, str]:
+    if not url:
+        raise ValueError('url is empty')
+
+    if url.startswith('tos://'):
+        url = url[len('tos://'):]
+
+    if not url:
+        raise ValueError('bucket is empty')
+
+    url = url.split('/', maxsplit=1)
+    if len(url) == 1:
+        bucket = url[0]
+        prefix = ''
+    else:
+        bucket = url[0]
+        prefix = url[1]
+
+    if not bucket:
+        raise ValueError('bucket is empty')
+    return bucket, prefix
+
+
+def default_trans(obj: TosObjectReader) -> TosObjectReader:
+    return obj
+
+
+def gen_dataset_from_urls(urls: Union[str, Iterator[str]], _: TosClient) -> Iterator[TosObjectMeta]:
+    if isinstance(urls, str):
+        urls = [urls]
+    return (TosObjectMeta(bucket, key) for bucket, key in [parse_tos_url(url) for url in urls])
+
+
+def gen_dataset_from_prefix(prefix: str, client: TosClient) -> Iterator[TosObjectMeta]:
+    bucket, prefix = parse_tos_url(prefix)
+    return iter(TosObjectIterator(bucket, prefix, client))

tostorchconnector-1.0.1/tostorchconnector/tos_error.py

@@ -0,0 +1,11 @@
+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

tostorchconnector-1.0.1/tostorchconnector/tos_iterable_dataset.py

@@ -0,0 +1,99 @@
+import logging
+from functools import partial
+from typing import Iterator, Any, Optional, Callable, Union
+
+import torch
+
+from . import TosObjectReader
+from .tos_client import CredentialProvider, TosClientConfig, TosClient, TosLogConfig
+from .tos_common import default_trans, gen_dataset_from_urls, gen_dataset_from_prefix
+from .tos_object_meta import TosObjectMeta
+
+log = logging.getLogger(__name__)
+
+
+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,
+                 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
+        self._gen_dataset = gen_dataset
+        self._region = region
+        self._endpoint = endpoint
+        self._trans = trans
+        self._cred = cred
+        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
+
+    @classmethod
+    def from_urls(cls, urls: Union[str, Iterator[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=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)
+
+    @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=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)
+
+    def __iter__(self) -> Iterator[Any]:
+        worker_id = 0
+        num_workers = 1
+        if self._sharding:
+            worker_info = torch.utils.data.get_worker_info()
+            if worker_info is not None:
+                worker_id = worker_info.id
+                num_workers = worker_info.num_workers
+
+        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()),
+            )
+
+        part_dataset = (
+            obj
+            for idx, obj in enumerate(self._gen_dataset(self._get_tos_client()))
+            if idx % self._world_size == self._rank
+        )
+
+        part_dataset = (
+            obj
+            for idx, obj in enumerate(part_dataset)
+            if idx % num_workers == worker_id
+        )
+        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)
+        return self._trans(obj)

tostorchconnector-1.0.1/tostorchconnector/tos_map_dataset.py

@@ -0,0 +1,78 @@
+import logging
+from functools import partial
+from typing import Any, Callable, Iterator, Optional, List, Union
+
+import torch
+
+from . import TosObjectReader
+from .tos_client import CredentialProvider, TosClientConfig, TosClient, TosLogConfig
+from .tos_common import default_trans, gen_dataset_from_prefix, \
+    gen_dataset_from_urls
+from .tos_object_meta import TosObjectMeta
+
+log = logging.getLogger(__name__)
+
+
+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,
+                 cred: Optional[CredentialProvider] = None,
+                 client_conf: Optional[TosClientConfig] = None,
+                 log_conf: Optional[TosLogConfig] = None, use_native_client=True):
+        self._client: Optional[TosClient] = None
+        self._gen_dataset = gen_dataset
+        self._region = region
+        self._endpoint = endpoint
+        self._trans = trans
+        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
+
+    @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.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)
+
+    @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.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)
+
+    def __getitem__(self, i: int) -> Any:
+        return self._trans_tos_object(i)
+
+    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):
+        if self._dataset is None:
+            self._dataset = list(self._gen_dataset(self._get_tos_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)
+        return self._trans(obj)

tostorchconnector-1.0.1/tostorchconnector/tos_object_meta.py

@@ -0,0 +1,25 @@
+from typing import Optional
+
+
+class TosObjectMeta(object):
+    def __init__(self, bucket: str, key: str, size: Optional[int] = None, etag: Optional[str] = None):
+        self._bucket = bucket
+        self._key = key
+        self._size = size
+        self._etag = etag
+
+    @property
+    def bucket(self) -> str:
+        return self._bucket
+
+    @property
+    def key(self) -> str:
+        return self._key
+
+    @property
+    def size(self) -> Optional[int]:
+        return self._size
+
+    @property
+    def etag(self) -> Optional[str]:
+        return self._etag

tostorchconnector-1.0.1/tostorchconnector/tos_object_reader.py

@@ -0,0 +1,170 @@
+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 tos.models2 import GetObjectOutput
+from tosnativeclient.tosnativeclient import ReadStream
+
+from .tos_object_meta import TosObjectMeta
+
+log = logging.getLogger(__name__)
+
+
+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]):
+        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_offset = 0
+        self._total_size: Optional[int] = None
+        self._read_offset = 0
+        self._buffer = io.BytesIO()
+
+    @property
+    def bucket(self) -> str:
+        return self._bucket
+
+    @property
+    def key(self) -> str:
+        return self._key
+
+    def read(self, size: Optional[int] = None) -> Optional[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:
+            # 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:
+                            break
+                        self._object_stream_offset += len(chunk)
+                        self._buffer.write(chunk)
+                finally:
+                    self._object_stream.close()
+            else:
+                for chunk in self._object_stream:
+                    self._buffer.write(chunk)
+
+            self._total_size = self._buffer.tell()
+        else:
+            self.seek(size, SEEK_CUR)
+
+        self._buffer.seek(current_read_offset)
+        data = self._buffer.read(size)
+        self._read_offset = self._buffer.tell()
+        return data
+
+    def readinto(self, buf) -> Optional[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)
+        readed = self._buffer.readinto(buf)
+        self._read_offset = self._buffer.tell()
+        return readed
+
+    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')
+
+        if offset < 0:
+            raise ValueError(f'invalid seek offset {offset}')
+
+        if offset > self._buffer_size():
+            self._prefetch_to_offset(offset)
+
+        self._read_offset = self._buffer.seek(offset)
+        return self._read_offset
+
+    def tell(self) -> int:
+        return self._read_offset
+
+    def readable(self):
+        return True
+
+    def writable(self):
+        return False
+
+    def seekable(self):
+        return True
+
+    @cached_property
+    def _object_meta(self) -> TosObjectMeta:
+        return self._get_object_meta()
+
+    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 _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_meta.size
+        return self._total_size
+
+    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:
+                        break
+                    size += self._buffer.write(chunk)
+                    self._object_stream_offset += len(chunk)
+                self._total_size = self._buffer.tell()
+            finally:
+                self._object_stream.close()
+        else:
+            try:
+                while offset > size:
+                    size += self._buffer.write(next(self._object_stream))
+            except StopIteration:
+                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

tostorchconnector-1.0.1/tostorchconnector/tos_object_writer.py

@@ -0,0 +1,83 @@
+import io
+import logging
+import threading
+from typing import Callable, Optional, Any
+
+from tos.models2 import PutObjectOutput
+from tosnativeclient.tosnativeclient import WriteStream
+
+log = logging.getLogger(__name__)
+
+
+class PutObjectStream(object):
+    def __init__(self, put_object: Callable[[io.BytesIO], PutObjectOutput]):
+        self._put_object = put_object
+        self._buffer = io.BytesIO()
+
+    def write(self, data):
+        self._buffer.write(data)
+
+    def close(self):
+        self._buffer.seek(0)
+        _ = self._put_object(self._buffer)
+
+
+class TosObjectWriter(io.BufferedIOBase):
+
+    def __init__(self, bucket: str, key: str, put_object_stream: WriteStream | PutObjectStream):
+        if not bucket:
+            raise ValueError('bucket is empty')
+        self._bucket = bucket
+        self._key = key
+        self._put_object_stream = put_object_stream
+        self._write_offset = 0
+        self._closed = False
+        self._lock = threading.Lock()
+
+    @property
+    def bucket(self) -> str:
+        return self._bucket
+
+    @property
+    def key(self) -> str:
+        return self._key
+
+    def __enter__(self):
+        self._write_offset = 0
+        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:
+            self.close()
+
+    def write(self, data) -> int:
+        self._put_object_stream.write(data)
+        self._write_offset += len(data)
+        return len(data)
+
+    def close(self) -> None:
+        with self._lock:
+            if self._closed:
+                return
+            self._closed = True
+            self._put_object_stream.close()
+
+    def tell(self) -> int:
+        return self._write_offset
+
+    def flush(self) -> None:
+        pass
+
+    def readable(self):
+        return False
+
+    def writable(self):
+        return True
+
+    def seekable(self):
+        return False

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

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

tostorchconnector-1.0.1/tostorchconnector.egg-info/SOURCES.txt

@@ -0,0 +1,20 @@
+LICENSE
+README.md
+pyproject.toml
+tests/test_tos_dataset.py
+tests/test_tosnativeclient.py
+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/tos_object_reader.py
+tostorchconnector/tos_object_writer.py
+tostorchconnector.egg-info/PKG-INFO
+tostorchconnector.egg-info/SOURCES.txt
+tostorchconnector.egg-info/dependency_links.txt
+tostorchconnector.egg-info/requires.txt
+tostorchconnector.egg-info/top_level.txt

tostorchconnector-1.0.1/tostorchconnector.egg-info/top_level.txt

@@ -0,0 +1 @@
+tostorchconnector

tostorchconnector-1.0.0/tostorchconnector/tostorchconnector.egg-info/SOURCES.txt

@@ -1,10 +0,0 @@
-LICENSE
-README.md
-pyproject.toml
-tests/test_tos_dataset.py
-tests/test_tosnativeclient.py
-tostorchconnector/tostorchconnector.egg-info/PKG-INFO
-tostorchconnector/tostorchconnector.egg-info/SOURCES.txt
-tostorchconnector/tostorchconnector.egg-info/dependency_links.txt
-tostorchconnector/tostorchconnector.egg-info/requires.txt
-tostorchconnector/tostorchconnector.egg-info/top_level.txt

tostorchconnector-1.0.0/tostorchconnector/tostorchconnector.egg-info/top_level.txt

@@ -1 +0,0 @@
-