This is an automated email from the ASF dual-hosted git repository.
JingsongLi pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/paimon.git
The following commit(s) were added to refs/heads/master by this push:
new 219f9ff9fe [python] Fix local cache being disabled in spawned workers
(#10479)
219f9ff9fe is described below
commit 219f9ff9fe861759689978ccc1e0d6d4c94e12e6
Author: XiaoHongbo <[email protected]>
AuthorDate: Sun Oct 11 22:58:25 2026 +0800
[python] Fix local cache being disabled in spawned workers (#10479)
---
docs/docs/program-api/file-cache.mdx | 7 +-
.../pypaimon/filesystem/caching_file_io.py | 78 ++++++-
.../pypaimon/tests/caching_file_io_test.py | 252 ++++++++++++++++++++-
3 files changed, 320 insertions(+), 17 deletions(-)
diff --git a/docs/docs/program-api/file-cache.mdx
b/docs/docs/program-api/file-cache.mdx
index 6d4b2fd24a..67320f7a17 100644
--- a/docs/docs/program-api/file-cache.mdx
+++ b/docs/docs/program-api/file-cache.mdx
@@ -167,6 +167,7 @@ reference to the shared manager; the last reference
releases the shared entry. A
### PyPaimon workers
-PyPaimon's cache is removed when `CachingFileIO` is pickled. An unpickled copy
reads directly from
-storage. Create a catalog with cache options on each Python worker when you
need caching there.
-The Java worker lifecycle above does not apply to Python serialization.
+With `spawn`, PyPaimon preserves cache configuration when `CachingFileIO` is
pickled. Workers reuse
+process-local managers across tasks; memory entries are not copied from the
parent, and disk caches
+scan their directory on first restore. With `fork`, create the catalog after
forking to avoid
+inheriting cache locks.
diff --git a/paimon-python/pypaimon/filesystem/caching_file_io.py
b/paimon-python/pypaimon/filesystem/caching_file_io.py
index 1ddbb76ee0..123e805429 100644
--- a/paimon-python/pypaimon/filesystem/caching_file_io.py
+++ b/paimon-python/pypaimon/filesystem/caching_file_io.py
@@ -28,17 +28,76 @@ the delegate FileIO.
import hashlib
import os
import threading
+import uuid
+import weakref
from collections import OrderedDict
+from concurrent.futures import Future
from typing import Optional
from pypaimon.common.file_io import FileIO, supports_pread, pread
from pypaimon.utils.file_type import FileType
+class _WorkerCacheRegistry:
+ """Share active managers and retain a few after their task completes."""
+
+ def __init__(self):
+ self._lock = threading.Lock()
+ self._active = weakref.WeakValueDictionary()
+ self._recent = OrderedDict()
+ self._pending = {}
+
+ def _remember(self, key, manager):
+ self._recent[key] = manager
+ self._recent.move_to_end(key)
+ if len(self._recent) > 8:
+ self._recent.popitem(last=False)
+
+ def restore(self, identity, cache_type, args):
+ key = (os.getpid(), identity, cache_type, args)
+ with self._lock:
+ manager = self._active.get(key)
+ if manager is not None:
+ self._remember(key, manager)
+ return manager
+ pending = self._pending.get(key)
+ create = pending is None
+ if create:
+ pending = Future()
+ self._pending[key] = pending
+
+ if not create:
+ return pending.result()
+
+ try:
+ manager = cache_type(*args)
+ manager._identity = identity
+ except BaseException as error:
+ with self._lock:
+ del self._pending[key]
+ pending.set_exception(error)
+ raise
+
+ with self._lock:
+ self._active[key] = manager
+ self._remember(key, manager)
+ del self._pending[key]
+ pending.set_result(manager)
+ return manager
+
+
+_WORKER_CACHE_REGISTRY = _WorkerCacheRegistry()
+
+
+def _restore_cache_manager(identity, cache_type, args):
+ return _WORKER_CACHE_REGISTRY.restore(identity, cache_type, args)
+
+
class LocalMemoryCacheManager:
"""Block-level in-memory cache with LRU eviction."""
def __init__(self, max_size_bytes: int, block_size: int = 1 * 1024 * 1024):
+ self._identity = uuid.uuid4().hex
self._max_size_bytes = max_size_bytes
self._block_size = block_size
self._lock = threading.Lock()
@@ -46,6 +105,11 @@ class LocalMemoryCacheManager:
self._cache: OrderedDict = OrderedDict()
self._file_size_cache: dict = {}
+ def __reduce__(self):
+ # Rebuild process-local state without serializing cached blocks or
locks.
+ return (_restore_cache_manager,
+ (self._identity, type(self), (self._max_size_bytes,
self._block_size)))
+
@property
def block_size(self) -> int:
return self._block_size
@@ -83,6 +147,7 @@ class LocalDiskCacheManager:
def __init__(self, cache_dir: str, max_size_bytes: int,
block_size: int = 1 * 1024 * 1024):
+ self._identity = uuid.uuid4().hex
self._cache_dir = cache_dir
self._max_size_bytes = max_size_bytes
self._block_size = block_size
@@ -94,6 +159,11 @@ class LocalDiskCacheManager:
os.makedirs(cache_dir, exist_ok=True)
self._current_size = self._scan_and_populate_index()
+ def __reduce__(self):
+ return (_restore_cache_manager,
+ (self._identity, type(self),
+ (self._cache_dir, self._max_size_bytes, self._block_size)))
+
@property
def block_size(self) -> int:
return self._block_size
@@ -329,8 +399,7 @@ class CachingFileIO(FileIO):
"""FileIO wrapper that caches reads at block granularity.
Only file types in the whitelist are cached. Others are read directly
- from the delegate. After pickling/unpickling, the cache is None and reads
- fall through to the delegate directly.
+ from the delegate. Cache managers are restored process-locally when
unpickled.
"""
def __init__(self, delegate: FileIO, cache, whitelist=None):
@@ -342,9 +411,8 @@ class CachingFileIO(FileIO):
self._whitelist = whitelist
def __getstate__(self):
- state = self.__dict__.copy()
- state['_cache'] = None
- return state
+ # Python 3.10 otherwise finds the delegate's hook through __getattr__.
+ return self.__dict__.copy()
# Fallback caps when local-cache.max-size is unset (memory shares the
heap).
_DEFAULT_MEMORY_CACHE_MAX_SIZE = 256 * 1024 * 1024
diff --git a/paimon-python/pypaimon/tests/caching_file_io_test.py
b/paimon-python/pypaimon/tests/caching_file_io_test.py
index 940a4f4ab2..88edfbcf4b 100644
--- a/paimon-python/pypaimon/tests/caching_file_io_test.py
+++ b/paimon-python/pypaimon/tests/caching_file_io_test.py
@@ -18,20 +18,66 @@
"""Tests for LocalDiskCacheManager, CachingInputStream, and CachingFileIO."""
+import gc
import io
+import multiprocessing
import os
import pickle
import shutil
import tempfile
import threading
import unittest
-from unittest.mock import MagicMock
+from concurrent.futures import ThreadPoolExecutor
+from unittest.mock import MagicMock, patch
from pypaimon.filesystem.caching_file_io import (
LocalDiskCacheManager,
+ LocalMemoryCacheManager,
CachingFileIO,
CachingInputStream,
)
+from pypaimon.filesystem.local_file_io import LocalFileIO
+
+
+class _CountingLocalFileIO(LocalFileIO):
+
+ def __init__(self):
+ super().__init__()
+ self.opens = 0
+
+ def new_input_stream(self, path):
+ self.opens += 1
+ return super().new_input_stream(path)
+
+
+def _read_cached_twice(file_io, path, connection):
+ values = []
+ for _ in range(2):
+ with file_io.new_input_stream(path) as stream:
+ values.append(stream.read())
+ connection.send((values, file_io._delegate.opens))
+ connection.close()
+
+
+def _read_independent_pickles(payload, path, connection):
+ first = pickle.loads(payload)
+ second = pickle.loads(payload)
+ values = []
+ for file_io in (first, second):
+ with file_io.new_input_stream(path) as stream:
+ values.append(stream.read())
+ connection.send((values, first._delegate.opens, second._delegate.opens,
+ first._cache is second._cache))
+ connection.close()
+
+
+def _read_rest_cached_twice(file_io, path, connection):
+ values = []
+ for _ in range(2):
+ with file_io.new_input_stream(path) as stream:
+ values.append(stream.read())
+ connection.send((values, file_io._cache.get_block(path, 0)))
+ connection.close()
class LocalDiskCacheManagerTest(unittest.TestCase):
@@ -386,20 +432,208 @@ class CachingFileIOTest(unittest.TestCase):
caching_io.new_output_stream("/out")
delegate.new_output_stream.assert_called_once_with("/out")
- def test_pickle_drops_cache(self):
- from pypaimon.filesystem.local_file_io import LocalFileIO
+ def test_pickle_preserves_memory_cache_configuration(self):
from pypaimon.utils.file_type import FileType
- cache = LocalDiskCacheManager(
- self.cache_dir, 2 ** 63 - 1, block_size=64)
+ cache = LocalMemoryCacheManager(128, block_size=64)
+ cache.put_block("snapshot-1", 0, b"parent data")
+ cache.put_file_size("snapshot-1", 11)
caching_io = CachingFileIO(
- LocalFileIO(), cache, {FileType.META, FileType.DATA})
-
+ _CountingLocalFileIO(), cache, {FileType.META})
restored = pickle.loads(pickle.dumps(caching_io))
+ self.assertIsInstance(restored._cache, LocalMemoryCacheManager)
+ self.assertEqual(128, restored._cache._max_size_bytes)
+ self.assertEqual(64, restored._cache.block_size)
+ self.assertIsNone(restored._cache.get_block("snapshot-1", 0))
+ self.assertEqual(-1, restored._cache.get_file_size("snapshot-1"))
+ self.assertEqual({FileType.META}, restored._whitelist)
+ self.assertEqual(b"parent data", cache.get_block("snapshot-1", 0))
+ self.assertIsNot(cache._lock, restored._cache._lock)
+
+ def test_pickle_preserves_disk_cache_configuration(self):
+ cache = LocalDiskCacheManager(self.cache_dir, 128, block_size=64)
+ cache.put_block("snapshot-1", 0, b"persisted block")
+ restored = pickle.loads(pickle.dumps(
+ CachingFileIO(_CountingLocalFileIO(), cache)))
+ self.assertIsInstance(restored._cache, LocalDiskCacheManager)
+ self.assertEqual(self.cache_dir, restored._cache._cache_dir)
+ self.assertEqual(128, restored._cache._max_size_bytes)
+ self.assertEqual(64, restored._cache.block_size)
+ self.assertEqual(b"persisted block",
+ restored._cache.get_block("snapshot-1", 0))
+ self.assertIsNot(cache._lock, restored._cache._lock)
+
+ def test_pickle_preserves_shared_cache(self):
+ cache = LocalMemoryCacheManager(128, 64)
+ first, second = pickle.loads(pickle.dumps((
+ CachingFileIO(_CountingLocalFileIO(), cache),
+ CachingFileIO(_CountingLocalFileIO(), cache))))
+ self.assertIs(first._cache, second._cache)
+ first._cache.put_block("snapshot-1", 0, b"worker block")
+ self.assertEqual(b"worker block",
+ second._cache.get_block("snapshot-1", 0))
+
+ def test_independent_pickles_reuse_worker_cache(self):
+ for disk in (False, True):
+ with self.subTest(disk=disk):
+ cache = (LocalDiskCacheManager(
+ os.path.join(self.cache_dir, "blocks"), 128, 64)
+ if disk else LocalMemoryCacheManager(128, 64))
+ payload = pickle.dumps(CachingFileIO(_CountingLocalFileIO(),
cache))
+ first = pickle.loads(payload)
+ second = pickle.loads(payload)
+ self.assertIs(first._cache, second._cache)
+ self.assertIsNot(cache, first._cache)
+ first._cache.put_block("snapshot-1", 0, b"worker block")
+ self.assertEqual(b"worker block",
+ second._cache.get_block("snapshot-1", 0))
+
+ def test_separate_memory_cache_identities_remain_isolated(self):
+ first = pickle.loads(pickle.dumps(LocalMemoryCacheManager(128, 64)))
+ second = pickle.loads(pickle.dumps(LocalMemoryCacheManager(128, 64)))
+ self.assertIsNot(first, second)
+ first.put_block("snapshot-1", 0, b"first catalog")
+ self.assertIsNone(second.get_block("snapshot-1", 0))
+
+ def test_active_manager_survives_recent_cache_eviction(self):
+ payload = pickle.dumps(LocalMemoryCacheManager(128, 64))
+ first = pickle.loads(payload)
+ first.put_block("snapshot-1", 0, b"worker block")
+ for _ in range(9):
+ pickle.loads(pickle.dumps(LocalMemoryCacheManager(128, 64)))
+ second = pickle.loads(payload)
+ self.assertIs(first, second)
+ self.assertEqual(b"worker block", second.get_block("snapshot-1", 0))
+
+ def test_recent_manager_survives_task_teardown(self):
+ payload = pickle.dumps(LocalMemoryCacheManager(128, 64))
+ first = pickle.loads(payload)
+ first.put_block("snapshot-1", 0, b"worker block")
+ del first
+ gc.collect()
+ second = pickle.loads(payload)
+ self.assertEqual(b"worker block", second.get_block("snapshot-1", 0))
+
+ def test_concurrent_restore_creates_one_manager(self):
+ payload = pickle.dumps(LocalMemoryCacheManager(128, 64))
+ barrier = threading.Barrier(8)
+ constructor_calls = []
+ original_init = LocalMemoryCacheManager.__init__
+
+ def slow_init(manager, *args, **kwargs):
+ constructor_calls.append(1)
+ threading.Event().wait(0.1)
+ original_init(manager, *args, **kwargs)
+
+ def restore():
+ barrier.wait(5)
+ return pickle.loads(payload)
+
+ with patch.object(LocalMemoryCacheManager, "__init__", slow_init):
+ with ThreadPoolExecutor(max_workers=8) as pool:
+ futures = [pool.submit(restore) for _ in range(8)]
+ restored = [future.result(timeout=10) for future in futures]
+ self.assertEqual(1, len(constructor_calls))
+ self.assertTrue(all(manager is restored[0] for manager in restored))
+
+ def test_pickle_disabled_cache(self):
+ restored = pickle.loads(pickle.dumps(
+ CachingFileIO(_CountingLocalFileIO(), None)))
self.assertIsNone(restored._cache)
- self.assertEqual({FileType.META, FileType.DATA}, restored._whitelist)
- self.assertIsInstance(restored._delegate, LocalFileIO)
+
+ def test_spawn_worker_reuses_cached_blocks(self):
+ path = os.path.join(self.cache_dir, "snapshot-1")
+ with open(path, "wb") as output:
+ output.write(b"cached data")
+ for disk in (False, True):
+ with self.subTest(disk=disk):
+ cache = (LocalDiskCacheManager(
+ os.path.join(self.cache_dir, "blocks"), 128, 64)
+ if disk else LocalMemoryCacheManager(128, 64))
+ caching_io = CachingFileIO(_CountingLocalFileIO(), cache)
+ context = multiprocessing.get_context("spawn")
+ parent, child = context.Pipe()
+ process = context.Process(
+ target=_read_cached_twice, args=(caching_io, path, child))
+ process.start()
+ child.close()
+ try:
+ self.assertTrue(parent.poll(30), "worker timed out")
+ values, opens = parent.recv()
+ process.join(10)
+ self.assertEqual(0, process.exitcode)
+ self.assertEqual([b"cached data", b"cached data"], values)
+ self.assertEqual(1, opens)
+ finally:
+ parent.close()
+ if process.is_alive():
+ process.terminate()
+ process.join(10)
+
+ def test_spawn_worker_reuses_cache_across_independent_pickles(self):
+ path = os.path.join(self.cache_dir, "snapshot-1")
+ with open(path, "wb") as output:
+ output.write(b"cached data")
+ cache = LocalMemoryCacheManager(128, 64)
+ payload = pickle.dumps(CachingFileIO(_CountingLocalFileIO(), cache))
+ context = multiprocessing.get_context("spawn")
+ parent, child = context.Pipe()
+ process = context.Process(
+ target=_read_independent_pickles, args=(payload, path, child))
+ process.start()
+ child.close()
+ try:
+ self.assertTrue(parent.poll(30), "worker timed out")
+ values, first_opens, second_opens, shared = parent.recv()
+ process.join(10)
+ self.assertEqual(0, process.exitcode)
+ self.assertEqual([b"cached data", b"cached data"], values)
+ self.assertEqual((1, 0), (first_opens, second_opens))
+ self.assertTrue(shared)
+ finally:
+ parent.close()
+ if process.is_alive():
+ process.terminate()
+ process.join(10)
+
+ def test_spawn_rest_file_io_preserves_cache_wrapper(self):
+ from pypaimon.catalog.rest.rest_token import RESTToken
+ from pypaimon.catalog.rest.rest_token_file_io import RESTTokenFileIO
+ from pypaimon.common.identifier import Identifier
+ from pypaimon.common.options import Options
+
+ path = os.path.join(self.cache_dir, "snapshot-1")
+ with open(path, "wb") as output:
+ output.write(b"cached data")
+ delegate = RESTTokenFileIO(
+ Identifier.create("db", "t"), "file://" + self.cache_dir,
+ Options({"uri": "http://127.0.0.1:1";, "token.provider": "bear",
+ "token": "local-test-token"}))
+ delegate.token = RESTToken({}, 2 ** 63 - 1)
+ file_io = CachingFileIO(delegate, LocalMemoryCacheManager(128, 64))
+ restored = pickle.loads(pickle.dumps(file_io))
+ self.assertIsInstance(restored._delegate, RESTTokenFileIO)
+ self.assertIsInstance(restored._cache, LocalMemoryCacheManager)
+
+ context = multiprocessing.get_context("spawn")
+ parent, child = context.Pipe()
+ process = context.Process(
+ target=_read_rest_cached_twice, args=(file_io, path, child))
+ process.start()
+ child.close()
+ try:
+ self.assertTrue(parent.poll(30), "worker timed out")
+ values, cached_block = parent.recv()
+ process.join(10)
+ self.assertEqual(0, process.exitcode)
+ self.assertEqual([b"cached data", b"cached data"], values)
+ self.assertEqual(b"cached data", cached_block)
+ finally:
+ parent.close()
+ if process.is_alive():
+ process.terminate()
+ process.join(10)
def test_to_filesystem_path_forwarded_to_delegate(self):
delegate = MagicMock()
