test_multiprocessing_safety.py passes

Author: softwareengineerprogrammer

src/geophires_x_client/__init__.py

@@ -1,4 +1,3 @@
-import json
 import os
 import sys
 import threading
@@ -7,18 +6,23 @@
 
 from geophires_x import GEOPHIRESv3 as geophires
 
+# Assuming these are in a sibling file or accessible path
 from .common import _get_logger
-from .geophires_input_parameters import EndUseOption
 from .geophires_input_parameters import GeophiresInputParameters
+from .geophires_input_parameters import ImmutableGeophiresInputParameters
 from .geophires_x_result import GeophiresXResult
 
 
 class GeophiresXClient:
+    """
+    A thread-safe and process-safe client for running GEOPHIRES simulations.
+    Relies on an explicit shutdown() call to clean up background processes.
+    """
+
     # --- Class-level shared resources ---
-    # These will be initialized lazily and shared across all instances and processes.
     _manager = None
     _cache = None
-    _lock = None  # This will be a process-safe RLock from the manager.
+    _lock = None
 
     # A standard threading lock to make the one-time initialization thread-safe.
     _init_lock = threading.Lock()
@@ -31,7 +35,6 @@ def __init__(self, enable_caching=True, logger_name=None):
         self._enable_caching = enable_caching
 
         # Lazy-initialize shared resources if they haven't been already.
-        # This approach is safe to call from multiple threads/processes.
         if GeophiresXClient._manager is None:
             self._initialize_shared_resources()
 
@@ -41,69 +44,63 @@ def _initialize_shared_resources(cls):
         Initializes the multiprocessing Manager and shared resources (cache, lock)
         in a thread-safe and process-safe manner.
         """
-        # Use a thread-safe lock to ensure this block only ever runs once
-        # across all threads in the main process.
         with cls._init_lock:
-            # The double-check locking pattern ensures we don't try to
-            # re-initialize if another thread finished while we were waiting.
             if cls._manager is None:
                 cls._manager = Manager()
                 cls._cache = cls._manager.dict()
-                cls._lock = cls._manager.RLock()  # The Manager now creates the lock.
+                cls._lock = cls._manager.RLock()
 
-    def get_geophires_result(self, input_params: GeophiresInputParameters) -> GeophiresXResult:
+    @classmethod
+    def shutdown(cls):
         """
-        Calculates a GEOPHIRES result in a thread-safe and process-safe manner.
+        Explicitly shuts down the background manager process.
+        This MUST be called when the application is finished with the client
+        to prevent orphaned processes.
         """
-        # Use the process-safe lock from the manager to make the check-then-act
-        # operation on the cache fully atomic across multiple processes.
-        with GeophiresXClient._lock:
-            cache_key = hash(input_params)
-            if self._enable_caching and cache_key in GeophiresXClient._cache:
-                return GeophiresXClient._cache[cache_key]
+        with cls._init_lock:
+            if cls._manager is not None:
+                cls._manager.shutdown()
+                cls._manager = None
+                cls._cache = None
+                cls._lock = None
 
-            # --- This section is now guaranteed to run only once per unique input ---
-            stash_cwd = Path.cwd()
-            stash_sys_argv = sys.argv
+    def get_geophires_result(self, input_params: GeophiresInputParameters) -> GeophiresXResult:
+        """
+        Calculates a GEOPHIRES result, using a cross-process cache to avoid
+        re-computing results for the same inputs. Caching is only effective
+        when providing an instance of ImmutableGeophiresInputParameters.
+        """
+        is_immutable = isinstance(input_params, ImmutableGeophiresInputParameters)
 
-            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 global state is restored even if geophires.main() fails
-                sys.argv = stash_sys_argv
-                os.chdir(stash_cwd)
+        if not (self._enable_caching and is_immutable):
+            return self._run_simulation(input_params)
 
-            self._logger.info(f'GEOPHIRES-X output file: {input_params.get_output_file_path()}')
+        cache_key = hash(input_params)
 
-            result = GeophiresXResult(input_params.get_output_file_path())
-            if self._enable_caching:
-                self._cache[cache_key] = result
+        with GeophiresXClient._lock:
+            if cache_key in GeophiresXClient._cache:
+                return GeophiresXClient._cache[cache_key]
 
+            result = self._run_simulation(input_params)
+            GeophiresXClient._cache[cache_key] = result
             return result
 
-
-if __name__ == '__main__':
-    # This block remains for direct testing of the script.
-    client = GeophiresXClient()
-    log = _get_logger()
-
-    params = GeophiresInputParameters(
-        {
-            'Print Output to Console': 0,
-            'End-Use Option': EndUseOption.DIRECT_USE_HEAT.value,
-            'Reservoir Model': 1,
-            'Time steps per year': 1,
-            'Reservoir Depth': 3,
-            'Gradient 1': 50,
-            'Maximum Temperature': 250,
-        }
-    )
-
-    result_ = client.get_geophires_result(params)
-    log.info(f'Breakeven price: ${result_.direct_use_heat_breakeven_price_USD_per_MMBTU}/MMBTU')
-    log.info(json.dumps(result_.result, indent=2))
+    def _run_simulation(self, input_params: GeophiresInputParameters) -> GeophiresXResult:
+        """Helper method to encapsulate the actual GEOPHIRES run."""
+        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())
+        return result

tests/geophires_x_client_tests/test_multiprocessing_safety.py

@@ -6,28 +6,24 @@
 from logging.handlers import QueueHandler
 from queue import Empty
 
-from geophires_x_client import EndUseOption
-
 # Important: We must be able to import the client and all parameter classes
 from geophires_x_client import GeophiresXClient
+from geophires_x_client.geophires_input_parameters import EndUseOption
 from geophires_x_client.geophires_input_parameters import ImmutableGeophiresInputParameters
 
 
-# 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.
+    This is the function that each worker process will execute.
+    It must be a top-level function to be picklable by multiprocessing.
     """
     # 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 initialization is now done in the worker, relying on the
-        # lazy-loading singleton pattern in the client itself.
+        # The client will use the Manager that was injected by the test's main process.
         client = GeophiresXClient(enable_caching=True)
         params = ImmutableGeophiresInputParameters(params_dict)
         result = client.get_geophires_result(params)
@@ -37,10 +33,8 @@ def run_client_in_process(params_dict: dict, log_queue: multiprocessing.Queue, r
 
 
 class TestMultiprocessingSafety(unittest.TestCase):
-    # By removing setUpClass and tearDownClass, we ensure each test is fully isolated.
-
     def setUp(self):
-        """Set up a shared set of parameters for each test."""
+        """Set up a unique set of parameters for each test."""
         self.params_dict = {
             'Print Output to Console': 0,
             'End-Use Option': EndUseOption.DIRECT_USE_HEAT.value,
@@ -54,25 +48,26 @@ def setUp(self):
     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.
+        processes and that the cache is shared between them. This test is now
+        fully self-contained to prevent resource conflicts with the test runner.
         """
         if sys.platform == 'win32':
             self.skipTest("The 'fork' multiprocessing context is not available on Windows.")
 
         ctx = multiprocessing.get_context('fork')
-        # THE FIX: Use the Manager as a context manager within the test.
-        # This guarantees it and all its resources (queues, etc.) are
-        # properly created and shut down for each individual test run.
+        # Use the Manager as a context manager. This is the key to ensuring
+        # all resources it creates (queues, etc.) are properly shut down
+        # at the end of the block, preventing deadlocks.
         with ctx.Manager() as manager:
-            log_queue = manager.Queue()
-            result_queue = manager.Queue()
-
-            # The client needs to be re-initialized inside the test to use the new manager.
-            # This is a bit of a workaround to reset the class-level singleton for the test.
+            # For this test to work, we MUST inject the test-specific manager
+            # into the client's class-level singleton attributes.
             GeophiresXClient._manager = manager
             GeophiresXClient._cache = manager.dict()
             GeophiresXClient._lock = manager.RLock()
 
+            log_queue = manager.Queue()
+            result_queue = manager.Queue()
+
             num_processes = 4
             process_timeout_seconds = 15
 
@@ -102,7 +97,7 @@ def test_client_runs_real_geophires_and_caches_across_processes(self):
             for p in processes:
                 p.join(timeout=process_timeout_seconds)
                 if p.is_alive():
-                    p.terminate()  # Forcefully end if stuck
+                    p.terminate()
                     self.fail(f'Process {p.pid} failed to terminate cleanly.')
 
             # --- Assertions ---
@@ -128,7 +123,7 @@ def test_client_runs_real_geophires_and_caches_across_processes(self):
                 f'\nTest passed: Detected {successful_runs} non-cached GEOPHIRES run(s) for {num_processes} requests.'
             )
 
-        # Reset the client's singleton state after the test to not interfere with others.
+        # CRITICAL: Reset the client's singleton state after the test to not interfere with other tests.
         GeophiresXClient._manager = None
         GeophiresXClient._cache = None
         GeophiresXClient._lock = None