diff --git a/tests/test_kv_storage_manager.py b/tests/test_kv_storage_manager.py index 7ac3744d..6f05136a 100644 --- a/tests/test_kv_storage_manager.py +++ b/tests/test_kv_storage_manager.py @@ -13,6 +13,7 @@ # See the License for the specific language governing permissions and # limitations under the License. import asyncio +import threading from unittest.mock import AsyncMock, MagicMock, patch import pytest @@ -476,3 +477,97 @@ def test_put_data_custom_backend_meta_length_mismatch_raises_error(test_data_for asyncio.run(manager.put_data(test_data_for_put_data["data"], test_data_for_put_data["metadata"])) assert "does not match" in str(exc_info.value) + + +@pytest.mark.asyncio +@patch.object(KVStorageManager, "_connect_to_controller", lambda self: None) +async def test_kv_backend_operations_do_not_block_event_loop(test_data_for_put_data): + mock_storage_client = MagicMock() + config = {"client_name": "MockClient"} + with patch(f"{STORAGE_CLIENT_FACTORY_PATH}.create", return_value=mock_storage_client): + manager = KVStorageManager(controller_info=MagicMock(), config=config) + manager.notify_data_update = AsyncMock() + + data = test_data_for_put_data["data"] + metadata = test_data_for_put_data["metadata"] + values = manager._generate_values(data) + loop_thread = threading.get_ident() + merge_thread = None + merge_tensors = manager._merge_tensors_to_tensordict + + def record_merge_thread(*args, **kwargs): + nonlocal merge_thread + merge_thread = threading.get_ident() + return merge_tensors(*args, **kwargs) + + manager._merge_tensors_to_tensordict = record_merge_thread + + operations = [ + ("put", lambda: manager.put_data(data, metadata), None), + ("get", lambda: manager.get_data(metadata), values), + ("clear", lambda: manager.clear_data(metadata), None), + ] + + try: + for name, invoke, result in operations: + started = threading.Event() + release = threading.Event() + call_thread = None + + def blocking_call( + *args, + _started=started, + _release=release, + _result=result, + **kwargs, + ): + nonlocal call_thread + call_thread = threading.get_ident() + _started.set() + _release.wait(timeout=2) + return _result + + getattr(mock_storage_client, name).side_effect = blocking_call + task = asyncio.create_task(invoke()) + for _ in range(100): + if started.is_set() or task.done(): + break + await asyncio.sleep(0.01) + + if task.done() and not started.is_set(): + await task + assert started.is_set(), f"storage_client.{name} was not called" + assert not task.done(), f"storage_client.{name} blocked the event loop" + assert call_thread != loop_thread + release.set() + await task + if name == "get": + assert merge_thread != loop_thread + finally: + manager.close() + + +@patch.object(KVStorageManager, "_connect_to_controller", lambda self: None) +def test_close_releases_kv_resources_once(): + mock_storage_client = MagicMock() + config = {"client_name": "MockClient"} + with patch(f"{STORAGE_CLIENT_FACTORY_PATH}.create", return_value=mock_storage_client): + manager = KVStorageManager(controller_info=MagicMock(), config=config) + + storage_executor = manager._get_storage_executor() + reconstruction_executor = manager._get_executor() + storage_executor.submit(lambda: None).result() + reconstruction_executor.submit(lambda: None).result() + assert storage_executor is not reconstruction_executor + + manager.close() + manager.close() + + assert manager._storage_executor is None + assert manager._multi_threads_executor is None + assert not manager._notify_thread.is_alive() + mock_storage_client.close.assert_called_once_with() + with pytest.raises(RuntimeError, match="cannot schedule new futures after shutdown"): + storage_executor.submit(lambda: None) + with pytest.raises(RuntimeError, match="cannot schedule new futures after shutdown"): + reconstruction_executor.submit(lambda: None) diff --git a/transfer_queue/storage/clients/base.py b/transfer_queue/storage/clients/base.py index 14b73e72..d44b6896 100644 --- a/transfer_queue/storage/clients/base.py +++ b/transfer_queue/storage/clients/base.py @@ -69,6 +69,10 @@ def clear(self, keys: list[str], custom_backend_meta=None) -> None: """Clear key-value pairs in the storage backend.""" raise NotImplementedError("Subclasses must implement clear") + def close(self) -> None: + """Release resources owned by the backend client, if any.""" + return None + class StorageClientFactory: """ diff --git a/transfer_queue/storage/managers/base.py b/transfer_queue/storage/managers/base.py index 9c32e213..af6b2358 100644 --- a/transfer_queue/storage/managers/base.py +++ b/transfer_queue/storage/managers/base.py @@ -14,12 +14,12 @@ # limitations under the License. import asyncio +import functools import inspect import itertools import os import threading import time -import weakref from abc import ABC, abstractmethod from concurrent.futures import ThreadPoolExecutor from typing import Any, Callable @@ -530,13 +530,14 @@ def __init__( the controller notify/handshake path, so they keep an independent context rather than drawing on a caller's shared socket budget. """ + self._closed = False client_name = config.get("client_name", None) if client_name is None: raise ValueError("Missing client_name in config") super().__init__(controller_info, config) self.storage_client = StorageClientFactory.create(client_name, config) + self._storage_executor: ThreadPoolExecutor | None = None self._multi_threads_executor: ThreadPoolExecutor | None = None - self._executor_finalizer = weakref.finalize(self, self._shutdown_executor, self._multi_threads_executor) @staticmethod def _generate_keys(field_names: list[str], global_indexes: list[int]) -> list[str]: @@ -579,41 +580,57 @@ def _generate_values(data: TensorDict) -> list[Any]: results.extend(field_data) return results - @staticmethod - def _shutdown_executor(thread_executor: ThreadPoolExecutor | None) -> None: - """ - A static method to ensure no strong reference to 'self' is held within the - finalizer's callback, enabling proper garbage collection. - """ - if thread_executor: - thread_executor.shutdown(wait=False) + def _get_num_threads(self) -> int: + """Bound per-manager thread pools according to the current Ray allocation.""" + if hasattr(self, "_num_threads"): + return self._num_threads + + ray_context = ray.get_runtime_context() + is_in_ray_actor_or_task = ray_context.get_actor_id() is not None or ray_context.get_task_id() is not None + + if is_in_ray_actor_or_task: + ray_assigned_cpus = ray_context.get_assigned_resources().get("CPU", 1) + num_threads = min(max(2, int(ray_assigned_cpus)), LIMIT_THREADS_PER_MANAGER_IN_RAY_ACTOR) + else: + num_threads = min(max(2, os.cpu_count() or 2), LIMIT_THREADS_PER_MANAGER_IN_DRIVER) + + self._num_threads = num_threads + return num_threads def _get_executor(self) -> ThreadPoolExecutor: - """Lazy Creating multi-thread executor for speeding up '_merge_tensors_to_tensordict'""" - if self._multi_threads_executor is None: - ray_context = ray.get_runtime_context() - is_in_ray_actor_or_task = ray_context.get_actor_id() is not None or ray_context.get_task_id() is not None - - if is_in_ray_actor_or_task: - # In ray actor: - ray_assigned_cpus = ray_context.get_assigned_resources().get("CPU", 1) - # num_threads must be 2 at least. - num_threads = max(2, int(ray_assigned_cpus)) - num_threads = min(num_threads, LIMIT_THREADS_PER_MANAGER_IN_RAY_ACTOR) - else: - # In Driver: - # num_threads must be 2 at least. - num_threads = max(2, os.cpu_count() or 2) - num_threads = min(num_threads, LIMIT_THREADS_PER_MANAGER_IN_DRIVER) + """Lazily create the executor used to reconstruct TensorDict fields.""" + if self._closed: + raise RuntimeError("KVStorageManager is closed") - self._num_threads = num_threads + if self._multi_threads_executor is None: self._multi_threads_executor = ThreadPoolExecutor( - max_workers=self._num_threads, thread_name_prefix="KVStorageManager" + max_workers=self._get_num_threads(), thread_name_prefix="KVStorageManager" ) assert self._multi_threads_executor is not None return self._multi_threads_executor + def _get_storage_executor(self) -> ThreadPoolExecutor: + """Lazily create the executor that owns synchronous backend calls. + + A separate pool lets get workers wait for parallel field reconstruction + without deadlocking concurrent gets through nested submissions. + """ + if self._closed: + raise RuntimeError("KVStorageManager is closed") + + if self._storage_executor is None: + self._storage_executor = ThreadPoolExecutor( + max_workers=self._get_num_threads(), thread_name_prefix="KVStorageManagerIO" + ) + + return self._storage_executor + + async def _run_storage_call(self, operation: Callable, *args, **kwargs): + """Run a synchronous backend operation without blocking the caller's event loop.""" + call = functools.partial(operation, *args, **kwargs) + return await asyncio.get_running_loop().run_in_executor(self._get_storage_executor(), call) + def _merge_tensors_to_tensordict(self, metadata: BatchMeta, values: list[Any]) -> TensorDict: """ Reconstruct a TensorDict from a list of values using metadata. @@ -747,8 +764,7 @@ async def put_data( keys = self._generate_keys(data_field_names, metadata.global_indexes) values = self._generate_values(data) - loop = asyncio.get_event_loop() - custom_backend_meta = await loop.run_in_executor(None, self.storage_client.put, keys, values) + custom_backend_meta = await self._run_storage_call(self.storage_client.put, keys, values) field_schema = extract_field_schema(data) @@ -795,10 +811,14 @@ async def get_data(self, metadata: BatchMeta) -> TensorDict: return TensorDict({}, batch_size=len(metadata)) keys = self._generate_keys(metadata.field_names, metadata.global_indexes) shapes, dtypes, custom_backend_meta = self._get_shape_type_custom_backend_meta_list(metadata) - values = self.storage_client.get( - keys=keys, shapes=shapes, dtypes=dtypes, custom_backend_meta=custom_backend_meta - ) - return self._merge_tensors_to_tensordict(metadata, values) + + def get_and_merge() -> TensorDict: + values = self.storage_client.get( + keys=keys, shapes=shapes, dtypes=dtypes, custom_backend_meta=custom_backend_meta + ) + return self._merge_tensors_to_tensordict(metadata, values) + + return await self._run_storage_call(get_and_merge) async def clear_data(self, metadata: BatchMeta) -> None: """Remove stored data associated with the given metadata.""" @@ -810,4 +830,34 @@ async def clear_data(self, metadata: BatchMeta) -> None: keys = self._generate_keys(metadata.field_names, metadata.global_indexes) _, _, custom_meta = self._get_shape_type_custom_backend_meta_list(metadata) - self.storage_client.clear(keys=keys, custom_backend_meta=custom_meta) + await self._run_storage_call(self.storage_client.clear, keys=keys, custom_backend_meta=custom_meta) + + def close(self) -> None: + """Wait for KV work, release backend resources, and stop manager services.""" + if self._closed: + return + self._closed = True + + storage_executor = getattr(self, "_storage_executor", None) + self._storage_executor = None + executor = getattr(self, "_multi_threads_executor", None) + self._multi_threads_executor = None + for executor_name, current_executor in ( + ("storage", storage_executor), + ("reconstruction", executor), + ): + if current_executor is None: + continue + try: + current_executor.shutdown(wait=True) + except Exception as e: + logger.warning(f"[{self.storage_manager_id}]: Error shutting down KV {executor_name} executor: {e}") + + storage_client = getattr(self, "storage_client", None) + if storage_client is not None: + try: + storage_client.close() + except Exception as e: + logger.warning(f"[{self.storage_manager_id}]: Error closing KV storage client: {e}") + + super().close()