WIP - progress on multiprocessing safety

Author: softwareengineerprogrammer

src/geophires_x_client/__init__.py

@@ -1,10 +1,12 @@
 import json
 import os
 import sys
+
+# --- MULTIPROCESSING CHANGES ---
 from multiprocessing import Manager
+from multiprocessing import RLock
 from pathlib import Path
 
-# noinspection PyPep8Naming
 from geophires_x import GEOPHIRESv3 as geophires
 
 from .common import _get_logger
@@ -14,63 +16,74 @@
 
 
 class GeophiresXClient:
-    # Use a multiprocessing Manager to create a cache and lock that are
-    # shared across all processes spawned by a ProcessPoolExecutor.
-    _manager = Manager()
-    _cache = _manager.dict()
-    _lock = _manager.Lock()
+    # --- LAZY-LOADED, PROCESS-SAFE SINGLETONS ---
+    # Define class-level placeholders. These will be shared across all instances.
+    _manager = None
+    _cache = None
+    _lock = RLock()  # Use a process-safe re-entrant lock
 
     def __init__(self, enable_caching=True, logger_name=None):
         if logger_name is None:
             logger_name = __name__
 
         self._logger = _get_logger(logger_name=logger_name)
-        self.enable_caching = enable_caching
+        self._enable_caching = enable_caching
+
+        # This method will safely initialize the shared manager and cache
+        # only when the first client instance is created.
+        self._initialize_shared_resources()
+
+    @classmethod
+    def _initialize_shared_resources(cls):
+        """
+        Initializes the multiprocessing Manager and shared cache dictionary.
+        This method is designed to be called safely by multiple processes,
+        ensuring the manager is only started once.
+        """
+        with cls._lock:
+            if cls._manager is None:
+                # This code is now protected. It won't run on module import.
+                # It runs only when the first GeophiresXClient is instantiated.
+                cls._manager = Manager()
+                cls._cache = cls._manager.dict()
 
     def get_geophires_result(self, input_params: GeophiresInputParameters) -> GeophiresXResult:
-        if not self.enable_caching:
-            return self._run_geophires(input_params)
-
-        cache_key = hash(input_params)
-        with self.__class__._lock:
-            # Use a lock to ensure the check-and-set operation is atomic.
-            if cache_key in self.__class__._cache:
-                # Cache hit
-                return self.__class__._cache[cache_key]
-
-            # If not in cache, we will run the simulation, still under the lock,
-            # to prevent a race condition with other processes.
-            # Cache miss
-            result = self._run_geophires(input_params)
-            self.__class__._cache[cache_key] = result
-            return result
-
-    def _run_geophires(self, input_params: GeophiresInputParameters) -> GeophiresXResult:
-        """Helper method to encapsulate the actual GEOPHIRES execution."""
-        stash_cwd = Path.cwd()
-        stash_sys_argv = sys.argv
+        # Use the class-level lock to protect access to the shared cache
+        # and the non-reentrant GEOPHIRES core.
+        with GeophiresXClient._lock:
+            cache_key = hash(input_params)
+            if self._enable_caching and cache_key in GeophiresXClient._cache:
+                return GeophiresXClient._cache[cache_key]
+
+            # ... (The rest of your logic remains the same)
+            stash_cwd = Path.cwd()
+            stash_sys_argv = sys.argv
+
+            sys.argv = ['', input_params.as_file_path(), input_params.get_output_file_path()]
+            try:
+                geophires.main(enable_geophires_logging_config=False)
+            except Exception as e:
+                raise RuntimeError(f'GEOPHIRES encountered an exception: {e!s}') from e
+            except SystemExit:
+                raise RuntimeError('GEOPHIRES exited without giving a reason') from None
+            finally:
+                sys.argv = stash_sys_argv
+                os.chdir(stash_cwd)
+
+            self._logger.info(f'GEOPHIRES-X output file: {input_params.get_output_file_path()}')
+
+            result = GeophiresXResult(input_params.get_output_file_path())
+            if self._enable_caching:
+                GeophiresXClient._cache[cache_key] = result
 
-        sys.argv = ['', input_params.as_file_path(), input_params.get_output_file_path()]
-        try:
-            geophires.main(enable_geophires_logging_config=False)
-        except Exception as e:
-            raise RuntimeError(f'GEOPHIRES encountered an exception: {e!s}') from e
-        except SystemExit:
-            raise RuntimeError('GEOPHIRES exited without giving a reason') from None
-        finally:
-            # Ensure we always restore the original state
-            sys.argv = stash_sys_argv
-            os.chdir(stash_cwd)
-
-        self._logger.info(f'GEOPHIRES-X output file: {input_params.get_output_file_path()}')
-        return GeophiresXResult(input_params.get_output_file_path())
+            return result
 
 
 if __name__ == '__main__':
+    # This block is safe, as it's protected from being run on import.
     client = GeophiresXClient()
     log = _get_logger()
 
-    # noinspection PyTypeChecker
     params = GeophiresInputParameters(
         {
             'Print Output to Console': 0,

tests/geophires_x_client_tests/test_multiprocessing_safety.py

@@ -0,0 +1,152 @@
+import logging
+import multiprocessing
+import sys
+import time
+import unittest
+from logging.handlers import QueueHandler
+from queue import Empty
+
+from geophires_x_client import EndUseOption
+from geophires_x_client import GeophiresInputParameters
+
+# Important: We must be able to import the client
+from geophires_x_client import GeophiresXClient
+
+
+# This is the function that each worker process will execute.
+# It must be a top-level function to be picklable by multiprocessing.
+def run_client_in_process(params_dict: dict, log_queue: multiprocessing.Queue, result_queue: multiprocessing.Queue):
+    """
+    Instantiates a client and runs a calculation, reporting results
+    and logs back to the main process via queues.
+    """
+    # Configure logging for this worker process to send messages to the shared queue.
+    root_logger = logging.getLogger()
+    root_logger.setLevel(logging.INFO)
+    root_logger.handlers = [QueueHandler(log_queue)]
+
+    try:
+        client = GeophiresXClient(enable_caching=True)
+        params = GeophiresInputParameters(params_dict)
+
+        # This now calls the REAL geophires.main via the client.
+        result = client.get_geophires_result(params)
+
+        # Put the primitive result into the queue to avoid serialization issues.
+        result_queue.put(result.direct_use_heat_breakeven_price_USD_per_MMBTU)
+    except Exception as e:
+        # Report any exceptions back to the main process.
+        result_queue.put(e)
+
+
+class TestMultiprocessingSafety(unittest.TestCase):
+    # Class-level attributes to manage shared resources across test runs.
+    _ctx = None
+    _client_for_setup = None
+
+    @classmethod
+    def setUpClass(cls):
+        """
+        Set up the multiprocessing context and start the shared Manager
+        process ONCE before any tests in this class run.
+        """
+        if sys.platform == 'win32':
+            # Skip all tests in this class if not on a fork-supporting OS.
+            raise unittest.SkipTest("The 'fork' multiprocessing context is not available on Windows.")
+
+        cls._ctx = multiprocessing.get_context('fork')
+        # Instantiating the client here creates the shared _manager and _cache
+        # that all child processes forked from this test will inherit.
+        cls._client_for_setup = GeophiresXClient()
+
+    @classmethod
+    def tearDownClass(cls):
+        """
+        Shut down the shared Manager process ONCE after all tests in this
+        class have finished. This is the key to preventing hanging processes.
+        """
+        if cls._client_for_setup and hasattr(cls._client_for_setup, '_manager'):
+            if cls._client_for_setup._manager is not None:
+                cls._client_for_setup._manager.shutdown()
+
+    def setUp(self):
+        """Set up a shared set of parameters for each test."""
+        # This setup runs before each individual test method.
+        self.params_dict = {
+            'Print Output to Console': 0,
+            'End-Use Option': EndUseOption.DIRECT_USE_HEAT.value,
+            'Reservoir Model': 1,
+            'Time steps per year': 1,
+            # Use nanoseconds to ensure each test run gets a unique cache key (Use a different value per run)
+            'Reservoir Depth': 4 + time.time_ns() / 1e19,
+            'Gradient 1': 50,
+            'Maximum Temperature': 550,
+        }
+
+    def test_client_runs_real_geophires_and_caches_across_processes(self):
+        """
+        Tests that GeophiresXClient can run the real geophires.main in multiple
+        processes and that the cache is shared between them.
+        """
+        log_queue = self._ctx.Queue()
+        result_queue = self._ctx.Queue()
+        num_processes = 4
+        # Timeout should be long enough for at least one successful run.
+        process_timeout_seconds = 5
+
+        processes = [
+            self._ctx.Process(target=run_client_in_process, args=(self.params_dict, log_queue, result_queue))
+            for _ in range(num_processes)
+        ]
+
+        for p in processes:
+            p.start()
+
+        # --- Robust Result Collection ---
+        results = []
+        for i in range(num_processes):
+            try:
+                result = result_queue.get(timeout=process_timeout_seconds)
+                results.append(result)
+            except Empty:
+                # Terminate running processes before failing to avoid hanging the suite
+                for p_cleanup in processes:
+                    if p_cleanup.is_alive():
+                        p_cleanup.terminate()
+                self.fail(f'Test timed out waiting for result #{i + 1}. A worker process likely crashed or is stuck.')
+
+        # --- Process Cleanup ---
+        # With the robust tearDownClass, a simple join is sufficient here.
+        for p in processes:
+            p.join(timeout=process_timeout_seconds)
+
+        # --- Assertions ---
+        # 1. Check that no process returned an exception.
+        for r in results:
+            self.assertNotIsInstance(r, Exception, f'A process failed with an exception: {r}')
+
+        # 2. Check that all processes got a valid, non-None result.
+        for r in results:
+            self.assertIsNotNone(r)
+            self.assertIsInstance(r, float)
+
+        # 3. CRITICAL: Assert that the expensive GEOPHIRES calculation was only run ONCE.
+        #    This assertion is expected to fail until the caching bug in the client is fixed.
+        log_records = []
+        while not log_queue.empty():
+            log_records.append(log_queue.get().getMessage())
+
+        cache_indicator_log = 'GEOPHIRES-X output file:'
+        successful_runs = sum(1 for record in log_records if cache_indicator_log in record)
+
+        self.assertEqual(
+            successful_runs,
+            1,
+            f'FAIL: GEOPHIRES was run {successful_runs} times instead of once, indicating the cross-process cache failed.',
+        )
+
+        print(f'\nDetected {successful_runs} non-cached GEOPHIRES run(s) for {num_processes} requests.')
+
+
+if __name__ == '__main__':
+    unittest.main()