perf(core): Explicitly set PyTorch intra-op threads.

src/ml_flashpoint/core/checkpoint_saver.py

@@ -36,7 +36,7 @@
 from ml_flashpoint.core.defaults import DIRTY_MARKER_SUFFIX, CheckpointFormat, default_metadata_object_name
 from ml_flashpoint.core.mlf_logging import get_logger
 from ml_flashpoint.core.tensor_header import TensorHeader
-from ml_flashpoint.core.utils import log_execution_time
+from ml_flashpoint.core.utils import get_accelerator_count, log_execution_time
 from ml_flashpoint.replication.replication_manager import ReplicationManager
 
 DEFAULT_INITIAL_BUFFER_SIZE_BYTES = 16 * 1000 * 1000 * 1000
@@ -426,67 +426,83 @@ def write_data(
         thread_count: int = 1,
     ) -> list[WriteResult]:
         thread_count = max(thread_count, 1)
+        num_cpus = os.cpu_count() or 1
+        num_ranks = max(get_accelerator_count(), 1)
+        # Use 50% of available CPU cores for PyTorch intra-op threads and evenly distribute them across ranks.
+        torch_thread_count = max(1, num_cpus // 2 // num_ranks // thread_count)
+        original_num_threads = torch.get_num_threads()
+        # Explicitly set PyTorch intra-op threads to optimize for performance.
+        # This also avoids potential runtime errors in tensor.copy_() with concurrent writers
+        torch.set_num_threads(torch_thread_count)
         _LOGGER.debug(
-            "%s starting multi-threaded write_data with thread_count: %d",
+            "%s starting multi-threaded write_data. thread_count: %d, original_num_threads: %d, "
+            "num_cpus: %d, num_ranks: %d, torch_thread_count: %d",
             self.__class__.__name__,
             thread_count,
+            original_num_threads,
+            num_cpus,
+            num_ranks,
+            torch_thread_count,
         )
+        try:
+            # Queue of ObjectWriteBuckets
+            object_items_queue: queue.Queue = queue.Queue()
+            for bucket in write_buckets:
+                object_items_queue.put(bucket)
+
+            # NOTE: There is support for multiple threads, to simplify modifying that setting, but we typically
+            # only use 1 thread.
+
+            results_from_threads: queue.Queue = queue.Queue()  # Queue for tuple[List[WriteResult], Exception]
+            threads = []
+
+            # Kick off additional threads to main thread, if any.
+            _LOGGER.debug("Spawning %d extra writer threads (in addition to the main thread).", thread_count - 1)
+            for i in range(1, thread_count):
+                thread = threading.Thread(
+                    target=self._write_to_buffer_from_queue_worker,
+                    args=(object_items_queue, results_from_threads, replicate_after_write, self._use_optimized_save),
+                    name=f"{self.__class__.__name__}-Thread-{i}",
+                )
+                threads.append(thread)
+                thread.start()
 
-        # Queue of ObjectWriteBuckets
-        object_items_queue: queue.Queue = queue.Queue()
-        for bucket in write_buckets:
-            object_items_queue.put(bucket)
-
-        # NOTE: There is support for multiple threads, to simplify modifying that setting, but we typically
-        # only use 1 thread.
-
-        results_from_threads: queue.Queue = queue.Queue()  # Queue for tuple[List[WriteResult], Exception]
-        threads = []
-
-        # Kick off additional threads to main thread, if any.
-        _LOGGER.debug("Spawning %d extra writer threads (in addition to the main thread).", thread_count - 1)
-        for i in range(1, thread_count):
-            thread = threading.Thread(
-                target=self._write_to_buffer_from_queue_worker,
-                args=(object_items_queue, results_from_threads, replicate_after_write, self._use_optimized_save),
-                name=f"{self.__class__.__name__}-Thread-{i}",
-            )
-            threads.append(thread)
-            thread.start()
-
-        # Main thread execution.
-        self._write_to_buffer_from_queue_worker(
-            object_items_queue, results_from_threads, replicate_after_write, self._use_optimized_save
-        )
-
-        for thread in threads:
-            thread.join()
-
-        all_results: list[WriteResult] = []
-        exceptions_raised: list[Exception] = []
-        # Collect all results, replication metadata, and exceptions
-        while not results_from_threads.empty():
-            try:
-                results, exception = results_from_threads.get_nowait()
-                if exception:
-                    exceptions_raised.append(exception)
-                elif results:
-                    all_results.extend(results)
-            except queue.Empty:
-                break
-
-        if exceptions_raised:
-            _LOGGER.error(
-                "'%s' encountered %d error(s) during multi-threaded write (will propagate the first one):\n%s.",
-                self.__class__.__name__,
-                len(exceptions_raised),
-                exceptions_raised,
+            # Main thread execution.
+            self._write_to_buffer_from_queue_worker(
+                object_items_queue, results_from_threads, replicate_after_write, self._use_optimized_save
             )
-            # Propagate the first exception encountered.
-            # TODO: propagate some combined exception, then update log msg (for now they are all logged above at least)
-            raise exceptions_raised[0]
 
-        return all_results
+            for thread in threads:
+                thread.join()
+
+            all_results: list[WriteResult] = []
+            exceptions_raised: list[Exception] = []
+            # Collect all results, replication metadata, and exceptions
+            while not results_from_threads.empty():
+                try:
+                    results, exception = results_from_threads.get_nowait()
+                    if exception:
+                        exceptions_raised.append(exception)
+                    elif results:
+                        all_results.extend(results)
+                except queue.Empty:
+                    break
+
+            if exceptions_raised:
+                _LOGGER.error(
+                    "'%s' encountered %d error(s) during multi-threaded write (will propagate the first one):\n%s.",
+                    self.__class__.__name__,
+                    len(exceptions_raised),
+                    exceptions_raised,
+                )
+                # Propagate the first exception encountered.
+                # TODO: propagate some combined exception, then update log msg
+                # (for now they are all logged above at least)
+                raise exceptions_raised[0]
+
+            return all_results
+        finally:
+            torch.set_num_threads(original_num_threads)
 
     @log_execution_time(logger=_LOGGER, name="async_replicate_object")
     def async_replicate_object(self, object_id: CheckpointObjectId) -> list[concurrent.futures.Future]:

src/ml_flashpoint/core/utils.py

@@ -33,6 +33,19 @@ def get_env_var_prefix() -> str:
     return "MLFLASHPOINT"
 
 
+def get_accelerator_count() -> int:
+    """Returns the number of accelerators available in the host.
+
+    Returns:
+        int: The number of accelerators.
+
+    """
+    if torch.cuda.is_available():
+        return torch.cuda.device_count()
+    else:
+        return 0
+
+
 def get_num_of_nodes() -> int:
     """
     Calculates the number of nodes in a distributed job without using LOCAL_WORLD_SIZE.
@@ -59,15 +72,7 @@ def get_num_of_nodes() -> int:
 
     world_size = dist.get_world_size()
 
-    if torch.cuda.is_available():
-        nprocs_per_node = torch.cuda.device_count()
-    else:
-        # This will fail for CPU-only distributed training.
-        raise RuntimeError(
-            "Cannot determine number of nodes for CPU-only training without "
-            "the `NNODES` or `LOCAL_WORLD_SIZE` environment variables."
-        )
-
+    nprocs_per_node = get_accelerator_count()
     if nprocs_per_node == 0:
         raise RuntimeError("torch.cuda.device_count() returned 0.")
 

tests/core/test_checkpoint_saver.py

@@ -106,6 +106,10 @@ def chkpt_object_manager(self):
     def replication_manager(self, mocker):
         return mocker.MagicMock(spec=ReplicationManager)
 
+    @pytest.fixture(autouse=True)
+    def mock_accelerator_count(self, mocker):
+        return mocker.patch("ml_flashpoint.core.checkpoint_saver.get_accelerator_count", return_value=1)
+
     @staticmethod
     def _tensor_write_data_for(tensor: torch.Tensor):
         return TensorWriteData(
@@ -213,6 +217,87 @@ def test_write_data_with_optimized_save_false(self, chkpt_object_manager, replic
         loaded_tensor = torch.load(data_io, weights_only=False)
         assert torch.equal(loaded_tensor, tensor_data)
 
+    @pytest.mark.parametrize(
+        "exception_in_worker",
+        [
+            None,
+            RuntimeError("Worker failed"),
+        ],
+    )
+    def test_write_data_resets_num_threads(
+        self, chkpt_object_manager, replication_manager, temp_dir_path, mocker, exception_in_worker
+    ):
+        # Given
+        saver = DefaultMLFlashpointCheckpointSaver(
+            global_rank_getter=lambda: 0,
+            local_rank_getter=lambda: 0,
+            global_barrier_func=lambda: None,
+            ckpt_obj_manager=chkpt_object_manager,
+            replication_manager=replication_manager,
+        )
+        checkpoint_id = CheckpointContainerId(os.path.join(temp_dir_path, "ckpt_threads"))
+
+        # Mock threading related calls
+        original_num_threads = 8
+        mocker.patch("torch.get_num_threads", return_value=original_num_threads)
+        mock_set_num_threads = mocker.patch("torch.set_num_threads")
+
+        # Mock worker to avoid actual writing and optionally raise exception
+        mock_worker = mocker.patch.object(saver, "_write_to_buffer_from_queue_worker")
+        if exception_in_worker:
+            mock_worker.side_effect = exception_in_worker
+
+        # When
+        if exception_in_worker:
+            with pytest.raises(RuntimeError, match="Worker failed"):
+                saver.write_data(checkpoint_id, [], replicate_after_write=False, thread_count=1)
+        else:
+            saver.write_data(checkpoint_id, [], replicate_after_write=False, thread_count=1)
+
+        # Then
+        # Verify it was reset to original_num_threads in finally block
+        assert mock_set_num_threads.call_args_list[-1] == mocker.call(original_num_threads)
+
+    def test_write_data_multithreaded(self, chkpt_object_manager, replication_manager, temp_dir_path):
+        # Given
+        saver = DefaultMLFlashpointCheckpointSaver(
+            global_rank_getter=lambda: 0,
+            local_rank_getter=lambda: 0,
+            global_barrier_func=lambda: None,
+            ckpt_obj_manager=chkpt_object_manager,
+            replication_manager=replication_manager,
+        )
+        checkpoint_id = CheckpointContainerId(os.path.join(temp_dir_path, "ckpt_threads_multi"))
+
+        # Create write items
+        num_items = 10
+        write_items = []
+        data_map = {}
+        for i in range(num_items):
+            tensor = torch.tensor([i], dtype=torch.int32)
+            index = MetadataIndex(fqn=f"item_{i}")
+            write_items.append(
+                WriteItem(
+                    index=index,
+                    type=WriteItemType.TENSOR,
+                    tensor_data=self._tensor_write_data_for(tensor),
+                )
+            )
+            data_map[index] = tensor
+
+        resolver = StubWriteItemResolver(data_map)
+
+        # Prepare buckets
+        buckets = saver.prepare_write_data(
+            checkpoint_id, write_items, resolver, object_name_prefix="data", bucket_count=4
+        )
+
+        # When
+        results = saver.write_data(checkpoint_id, buckets, replicate_after_write=False, thread_count=4)
+
+        # Then
+        assert len(results) == num_items
+
     @pytest.mark.parametrize(
         "checkpoint_id_suffix",
         [

tests/core/test_utils.py

@@ -118,7 +118,7 @@ def test_get_num_of_nodes_torch_only_cpu_only_training_raises_error(monkeypatch)
     monkeypatch.setattr(torch.distributed, "is_initialized", lambda: True)
     monkeypatch.setattr(torch.distributed, "get_world_size", lambda: 8)
     monkeypatch.setattr(torch.cuda, "is_available", lambda: False)
-    with pytest.raises(RuntimeError, match="Cannot determine number of nodes for CPU-only training"):
+    with pytest.raises(RuntimeError, match=re.escape("torch.cuda.device_count() returned 0.")):
         utils.get_num_of_nodes()
 
 
@@ -222,3 +222,26 @@ def test_get_env_val_int(self, env_vars, var_name, default_val, expected, monkey
 
         # Then
         assert result == expected
+
+
+class TestGetAcceleratorCount:
+    def test_get_accelerator_count_cuda_available(self, mocker):
+        # Given
+        mocker.patch("torch.cuda.is_available", return_value=True)
+        mocker.patch("torch.cuda.device_count", return_value=4)
+
+        # When
+        count = utils.get_accelerator_count()
+
+        # Then
+        assert count == 4
+
+    def test_get_accelerator_count_cuda_unavailable(self, mocker):
+        # Given
+        mocker.patch("torch.cuda.is_available", return_value=False)
+
+        # When
+        count = utils.get_accelerator_count()
+
+        # Then
+        assert count == 0