introduce lazy prefetch

Author: liyang

megfile/config.py

@@ -83,6 +83,7 @@ def set_log_level(level: T.Optional[T.Union[int, str]] = None):
 READER_MAX_BUFFER_SIZE = parse_quantity(
     os.getenv("MEGFILE_READER_MAX_BUFFER_SIZE") or 128 * 2**20
 )
+READER_LAZY_PREFETCH = parse_boolean(os.getenv("MEGFILE_READER_LAZY_PREFETCH"), False)
 
 # Multi-upload in aws s3 has a maximum of 10,000 parts,
 # so the maximum supported file size is MEGFILE_WRITE_BLOCK_SIZE * 10,000,

megfile/lib/base_prefetch_reader.py

@@ -82,9 +82,7 @@ def __init__(
 
         self._offset = 0
         self._cached_buffer = None
-        self._block_index = None  # Current block index
         self._seek_history = []
-
         self._seek_buffer(0)
 
         _logger.debug("open file: %r, mode: %s" % (self.name, self.mode))
@@ -98,7 +96,9 @@ def _futures(self) -> "LRUCacheFutureManager":
         return self._process_local("futures", self._get_futures)
 
     def _get_futures(self):
-        return LRUCacheFutureManager()
+        futures = LRUCacheFutureManager()
+        futures.register(self.name)
+        return futures
 
     @property
     @abstractmethod
@@ -207,9 +207,8 @@ def _read(self, size: int) -> bytes:
         if size == 0 or self._offset >= self._content_size:
             return b""
 
-        data = self._fetch_response(start=self._offset, end=self._offset + size - 1)[
-            "Body"
-        ].read()
+        resp = self._fetch_response(start=self._offset, end=self._offset + size - 1)
+        data = resp["Body"].read()
         self.seek(size, os.SEEK_CUR)
         return data
 
@@ -369,12 +368,17 @@ def _close(self):
 class LRUCacheFutureManager(OrderedDict):
     def __init__(self):
         super().__init__()
+        self._name = None
+
+    def register(self, name):
+        self._name = name
 
     def submit(self, executor, key, *args, **kwargs):
         if key in self:
             self.move_to_end(key, last=True)
             return
         self[key] = executor.submit(*args, **kwargs)
+        _logger.debug("submit future: %r, key: %r" % (self._name, key))
 
     @property
     def finished(self):
@@ -385,7 +389,12 @@ def result(self, key):
         return self[key].result()
 
     def cleanup(self, block_capacity: int):
+        keys = []
         while len(self) > block_capacity:
-            _, future = self.popitem(last=False)
+            key, future = self.popitem(last=False)
+            keys.append(key)
             if not future.done():
                 future.cancel()
+        if keys:
+            _logger.debug("cleanup futures: %r, keys: %s" % (self._name, keys))
+        return keys

megfile/lib/s3_prefetch_reader.py

@@ -4,6 +4,7 @@
 
 from megfile.config import (
     READER_BLOCK_SIZE,
+    READER_LAZY_PREFETCH,
     READER_MAX_BUFFER_SIZE,
     S3_MAX_RETRY_TIMES,
 )
@@ -62,7 +63,7 @@ def __init__(
         )
 
     def _get_content_size(self):
-        if self._block_capacity <= 0:
+        if self._block_capacity <= 0 or READER_LAZY_PREFETCH:
             response = self._client.head_object(Bucket=self._bucket, Key=self._key)
             self._content_etag = response.get("ETag")
             return int(response["ContentLength"])

megfile/lib/s3_share_cache_reader.py

@@ -101,16 +101,21 @@ def __init__(self):
         super().__init__()
         self._references = Counter()
 
-    def register(self, key):
-        self._references[key] += 1
-
-    def unregister(self, key):
-        self._references[key] -= 1
-        if self._references[key] == 0:
-            self._references.pop(key)
-            for key_tuple in list(self):
-                if key_tuple[0] != key:
+    def register(self, name):
+        self._references[name] += 1
+        _logger.debug("register reader: %r, count: %d" % (name, self._references[name]))
+
+    def unregister(self, name):
+        self._references[name] -= 1
+        _logger.debug(
+            "unregister reader: %r, count: %d" % (name, self._references[name])
+        )
+        if self._references[name] == 0:
+            self._references.pop(name)
+            for key in list(self):
+                if key[0] != name:
                     continue
-                future = self.pop(key_tuple)
+                future = self.pop(key)
                 if not future.done():
                     future.cancel()  # pragma: no cover
+            _logger.debug("cleanup all futures of reader: %r" % name)

megfile/s3_path.py

@@ -1000,6 +1000,7 @@ def s3_buffered_open(
         )
 
     if mode == "rb":
+        block_size = block_size or READER_BLOCK_SIZE
         if share_cache_key is not None:
             reader = S3ShareCacheReader(
                 bucket,
@@ -1008,7 +1009,7 @@ def s3_buffered_open(
                 s3_client=client,
                 max_retries=max_retries,
                 max_workers=max_workers,
-                block_size=block_size or READER_BLOCK_SIZE,
+                block_size=block_size,
                 block_forward=block_forward,
                 profile_name=s3_url._profile_name,
             )
@@ -1023,13 +1024,14 @@ def s3_buffered_open(
                 max_workers=max_workers,
                 max_buffer_size=max_buffer_size,
                 block_forward=block_forward,
-                block_size=block_size or READER_BLOCK_SIZE,
+                block_size=block_size,
                 profile_name=s3_url._profile_name,
             )
         if buffered or _is_pickle(reader):
-            reader = io.BufferedReader(reader)  # type: ignore
+            reader = io.BufferedReader(reader, buffer_size=block_size)  # type: ignore
         return reader
 
+    block_size = block_size or WRITER_BLOCK_SIZE
     if limited_seekable:
         if max_buffer_size is None:
             max_buffer_size = WRITER_MAX_BUFFER_SIZE
@@ -1038,7 +1040,7 @@ def s3_buffered_open(
             key,
             s3_client=client,
             max_workers=max_workers,
-            block_size=block_size or WRITER_BLOCK_SIZE,
+            block_size=block_size,
             max_buffer_size=max_buffer_size,
             profile_name=s3_url._profile_name,
         )
@@ -1050,12 +1052,12 @@ def s3_buffered_open(
             key,
             s3_client=client,
             max_workers=max_workers,
-            block_size=block_size or WRITER_BLOCK_SIZE,
+            block_size=block_size,
             max_buffer_size=max_buffer_size,
             profile_name=s3_url._profile_name,
         )
     if buffered or _is_pickle(writer):
-        writer = io.BufferedWriter(writer)  # type: ignore
+        writer = io.BufferedWriter(writer, buffer_size=block_size)  # type: ignore
     return writer
 
 

megfile/utils/__init__.py

@@ -18,6 +18,7 @@
 from threading import RLock
 from typing import IO, Callable, Optional
 
+from megfile.config import READER_LAZY_PREFETCH
 from megfile.utils.mutex import ProcessLocal, ThreadLocal
 
 
@@ -81,6 +82,9 @@ def is_writable(fileobj: IO) -> bool:
 
 def _is_pickle(fileobj) -> bool:
     """Test if File Object is pickle"""
+    if READER_LAZY_PREFETCH:
+        return False
+
     if fileobj.name.endswith(".pkl") or fileobj.name.endswith(".pickle"):
         return True