[None][fix] Restore tensor sharing strategy on errors

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(

tests/unittest/_torch/misc/test_share_tensor.py

@@ -1,7 +1,10 @@
 import multiprocessing as mp
 import unittest
+from unittest import mock
 
 import torch
+from torch.multiprocessing import (get_all_sharing_strategies,
+                                   get_sharing_strategy, set_sharing_strategy)
 
 from tensorrt_llm._torch.shared_tensor import SharedTensorContainer
 
@@ -88,6 +91,30 @@ 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_tensor_different_shapes(self):
         """Test CPU tensor sharing with different tensor shapes."""
         mp.set_start_method('spawn', force=True)