fix(mlops): use spawn for feature ingestion

sagemaker-mlops/src/sagemaker/mlops/feature_store/ingestion_manager_pandas.py

@@ -4,6 +4,7 @@
 
 import logging
 import math
+import multiprocessing
 import signal
 from concurrent.futures import ThreadPoolExecutor, as_completed
 from dataclasses import dataclass, field
@@ -25,6 +26,11 @@
 BATCH_WRITE_MAX_ENTRIES = 25
 
 
+def _initialize_worker():
+    """Ignore keyboard interrupts in child processes."""
+    signal.signal(signal.SIGINT, signal.SIG_IGN)
+
+
 class IngestionError(Exception):
     """Exception raised for errors during ingestion.
 
@@ -214,10 +220,11 @@ def _run_multi_process(
                 )
             )
 
-        def init_worker():
-            signal.signal(signal.SIGINT, signal.SIG_IGN)
-
-        self._processing_pool = Pool(self.max_processes, init_worker)
+        spawn_context = multiprocessing.get_context("spawn")
+        self._processing_pool = spawn_context.Pool(
+            processes=self.max_processes,
+            initializer=_initialize_worker,
+        )
 
         self._async_result = self._processing_pool.starmap_async(
             IngestionManagerPandas._run_multi_threaded,

sagemaker-mlops/tests/unit/sagemaker/mlops/feature_store/test_ingestion_manager_pandas.py

@@ -2,6 +2,8 @@
 # Licensed under the Apache License, Version 2.0
 """Unit tests for ingestion_manager_pandas.py"""
 
+import signal
+
 import pytest
 from unittest.mock import Mock, patch, MagicMock
 import pandas as pd
@@ -484,12 +486,12 @@ def test_multi_threaded_forwards_region(
         for call in mock_ingest.call_args_list:
             assert call[1]["region"] == "eu-west-1"
 
-    @patch("sagemaker.mlops.feature_store.ingestion_manager_pandas.Pool")
+    @patch("multiprocessing.get_context")
     def test_multi_process_args_include_region(
-        self, mock_pool_class, feature_definitions, sample_dataframe
+        self, mock_get_context, feature_definitions, sample_dataframe
     ):
         mock_pool = MagicMock()
-        mock_pool_class.return_value = mock_pool
+        mock_get_context.return_value.Pool.return_value = mock_pool
 
         manager = IngestionManagerPandas(
             feature_group_name="test-fg",
@@ -504,3 +506,59 @@ def test_multi_process_args_include_region(
         for process_args in starmap_args:
             # region is the last positional argument passed to _run_multi_threaded
             assert process_args[-1] == "eu-west-1"
+
+
+class TestIngestionManagerProcessPool:
+    """Regression tests for spawn-safe pandas ingestion pools."""
+
+    @patch("multiprocessing.get_context")
+    def test_multi_process_uses_spawn_context_and_module_initializer(self, mock_get_context):
+        spawn_context = Mock()
+        processing_pool = Mock()
+        async_result = Mock()
+        spawn_context.Pool.return_value = processing_pool
+        processing_pool.starmap_async.return_value = async_result
+        mock_get_context.return_value = spawn_context
+        manager = IngestionManagerPandas(
+            feature_group_name="test-fg",
+            feature_definitions={"id": {"FeatureType": "String"}},
+            max_workers=1,
+            max_processes=2,
+        )
+        data_frame = pd.DataFrame({"id": ["1", "2"]})
+
+        manager._run_multi_process(data_frame=data_frame, wait=False)
+
+        mock_get_context.assert_called_once_with("spawn")
+        spawn_context.Pool.assert_called_once()
+        pool_kwargs = spawn_context.Pool.call_args.kwargs
+        assert pool_kwargs["processes"] == 2
+        assert pool_kwargs["initializer"].__module__ == (
+            "sagemaker.mlops.feature_store.ingestion_manager_pandas"
+        )
+        assert pool_kwargs["initializer"].__qualname__ == "_initialize_worker"
+        with patch(
+            "sagemaker.mlops.feature_store.ingestion_manager_pandas.signal.signal"
+        ) as mock_signal:
+            pool_kwargs["initializer"]()
+        mock_signal.assert_called_once_with(signal.SIGINT, signal.SIG_IGN)
+        processing_pool.starmap_async.assert_called_once()
+        assert manager._async_result is async_result
+
+    def test_wait_closes_and_joins_spawn_pool(self):
+        manager = IngestionManagerPandas(
+            feature_group_name="test-fg",
+            feature_definitions={"id": {"FeatureType": "String"}},
+            max_workers=1,
+            max_processes=2,
+        )
+        manager._async_result = Mock()
+        manager._async_result.get.return_value = [[], []]
+        manager._processing_pool = Mock()
+
+        manager.wait(timeout=5)
+
+        manager._async_result.get.assert_called_once_with(timeout=5)
+        manager._processing_pool.close.assert_called_once_with()
+        manager._processing_pool.join.assert_called_once_with()
+        assert manager.failed_rows == []