diff --git a/tensorrt_llm/_torch/shared_tensor/shared_tensor.py b/tensorrt_llm/_torch/shared_tensor/shared_tensor.py index 24666511249d..1ed963fccd33 100644 --- a/tensorrt_llm/_torch/shared_tensor/shared_tensor.py +++ b/tensorrt_llm/_torch/shared_tensor/shared_tensor.py @@ -15,6 +15,8 @@ logger = logging.getLogger(__name__) +_sharing_strategy_lock = threading.Lock() + class _LocalTensorStore: """Process-local stash for same-process tensor handoff. @@ -452,14 +454,16 @@ def dump_to_dict(self) -> Dict[str, Any]: tensor_dict = SharedTensorContainer.cuda_handle_to_dict( self.tensor_handle) elif self.method_key == _SharedTensorRebuildMethodRegistry.REBUILD_CPU: - sharing_strategy = get_sharing_strategy() - # Here we use file_system sharing strategy to make it serializable between two non-python independent processes - set_sharing_strategy("file_system") storage = self.tensor_handle[1] meta_data = self.tensor_handle[2] - storage_handle = reduce_storage(storage) - # restore the original sharing strategy - set_sharing_strategy(sharing_strategy) + # Here we use file_system sharing strategy to make it serializable between two non-python independent processes + with _sharing_strategy_lock: + sharing_strategy = get_sharing_strategy() + set_sharing_strategy("file_system") + try: + storage_handle = reduce_storage(storage) + finally: + set_sharing_strategy(sharing_strategy) # exclude the first element which is the type of the storage storage_metadata = storage_handle[-1][1:] tensor_dict = SharedTensorContainer.cpu_handle_to_dict( diff --git a/tests/unittest/_torch/misc/test_share_tensor.py b/tests/unittest/_torch/misc/test_share_tensor.py index 548148255ebc..8eec65196075 100644 --- a/tests/unittest/_torch/misc/test_share_tensor.py +++ b/tests/unittest/_torch/misc/test_share_tensor.py @@ -1,7 +1,12 @@ import multiprocessing as mp +import threading import unittest +from unittest import mock import torch +from torch.multiprocessing import (get_all_sharing_strategies, + get_sharing_strategy, set_sharing_strategy) +from torch.multiprocessing.reductions import reduce_storage from tensorrt_llm._torch.shared_tensor import SharedTensorContainer @@ -88,6 +93,97 @@ def test_share_cpu_tensor(self): queue.close() queue.join_thread() + def test_share_cpu_tensor_restores_strategy_on_error(self): + """The process-wide sharing strategy is restored if serialization fails.""" + if "file_descriptor" not in get_all_sharing_strategies(): + self.skipTest("file_descriptor sharing is not available") + + original_strategy = get_sharing_strategy() + set_sharing_strategy("file_descriptor") + try: + container = SharedTensorContainer.from_tensor(self.ref_tensor) + + def fail_serialization(_storage): + self.assertEqual(get_sharing_strategy(), "file_system") + raise RuntimeError("serialization failed") + + with mock.patch( + "tensorrt_llm._torch.shared_tensor.shared_tensor.reduce_storage", + side_effect=fail_serialization): + with self.assertRaisesRegex(RuntimeError, + "serialization failed"): + container.dump_to_dict() + self.assertEqual(get_sharing_strategy(), "file_descriptor") + finally: + set_sharing_strategy(original_strategy) + + def test_share_cpu_tensor_serialization_is_serialized(self): + """Concurrent CPU serialization cannot overlap the strategy window.""" + if "file_descriptor" not in get_all_sharing_strategies(): + self.skipTest("file_descriptor sharing is not available") + + original_strategy = get_sharing_strategy() + first_reducer_entered = threading.Event() + release_first_reducer = threading.Event() + second_call_started = threading.Event() + second_reducer_entered = threading.Event() + call_count = 0 + call_count_lock = threading.Lock() + errors = [] + + def blocking_reduce(storage): + nonlocal call_count + with call_count_lock: + call_count += 1 + call_number = call_count + if call_number == 1: + first_reducer_entered.set() + if not release_first_reducer.wait(timeout=5): + raise TimeoutError("first reducer was not released") + else: + second_reducer_entered.set() + return reduce_storage(storage) + + def serialize(container, started=None): + if started is not None: + started.set() + try: + container.dump_to_dict() + except BaseException as error: + errors.append(error) + + set_sharing_strategy("file_descriptor") + first_thread = threading.Thread(target=serialize, + args=(SharedTensorContainer.from_tensor( + self.ref_tensor), )) + second_thread = threading.Thread( + target=serialize, + args=(SharedTensorContainer.from_tensor(self.ref_tensor), + second_call_started)) + try: + with mock.patch( + "tensorrt_llm._torch.shared_tensor.shared_tensor.reduce_storage", + side_effect=blocking_reduce): + first_thread.start() + self.assertTrue(first_reducer_entered.wait(timeout=5)) + second_thread.start() + self.assertTrue(second_call_started.wait(timeout=5)) + self.assertFalse(second_reducer_entered.wait(timeout=0.2)) + release_first_reducer.set() + first_thread.join(timeout=5) + second_thread.join(timeout=5) + + self.assertFalse(first_thread.is_alive()) + self.assertFalse(second_thread.is_alive()) + self.assertTrue(second_reducer_entered.is_set()) + self.assertEqual(errors, []) + self.assertEqual(get_sharing_strategy(), "file_descriptor") + finally: + release_first_reducer.set() + first_thread.join(timeout=5) + second_thread.join(timeout=5) + set_sharing_strategy(original_strategy) + def test_share_tensor_different_shapes(self): """Test CPU tensor sharing with different tensor shapes.""" mp.set_start_method('spawn', force=True)