chore: simplified parallel context

Author: ringoldsdev

laygo/context/parallel.py

@@ -3,131 +3,131 @@
 multiprocessing.Manager to share state across processes.
 """
 
+from collections.abc import Callable
 from collections.abc import Iterator
 import multiprocessing as mp
-from multiprocessing.managers import BaseManager
 from multiprocessing.managers import DictProxy
-from multiprocessing.synchronize import Lock
+from threading import Lock
 from typing import Any
+from typing import TypeVar
 
 from laygo.context.types import IContextHandle
 from laygo.context.types import IContextManager
 
-
-class _ParallelStateManager(BaseManager):
-  """A custom manager to expose a shared dictionary and lock."""
-
-  pass
+R = TypeVar("R")
 
 
 class ParallelContextHandle(IContextHandle):
   """
-  A lightweight, picklable "blueprint" for recreating a connection to the
-  shared context in a different process.
+  A lightweight, picklable handle that carries the actual shared objects
+  (the DictProxy and Lock) to worker processes.
   """
 
-  def __init__(self, address: tuple[str, int], manager_class: type["ParallelContextManager"]):
-    self.address = address
-    self.manager_class = manager_class
+  def __init__(self, shared_dict: DictProxy, lock: Lock):
+    self._shared_dict = shared_dict
+    self._lock = lock
 
   def create_proxy(self) -> "IContextManager":
     """
-    Creates a new instance of the ParallelContextManager in "proxy" mode
-    by initializing it with this handle.
+    Creates a new ParallelContextManager instance that wraps the shared
+    objects received by the worker process.
     """
-    return self.manager_class(handle=self)
+    return ParallelContextManager(handle=self)
 
 
 class ParallelContextManager(IContextManager):
   """
-  A context manager that uses a background multiprocessing.Manager to enable
-  state sharing across different processes.
-
-  This single class operates in two modes:
-  1. Server Mode (when created normally): It starts and manages the background
-     server process that holds the shared state.
-  2. Proxy Mode (when created with a handle): It acts as a client, connecting
-     to an existing server process to manipulate the shared state.
+  A context manager that enables state sharing across processes.
+
+  It operates in two modes:
+  1. Main Mode: When created normally, it starts a multiprocessing.Manager
+     and creates a shared dictionary and lock.
+  2. Proxy Mode: When created from a handle, it wraps a DictProxy and Lock
+     that were received from another process. It does not own the manager.
   """
 
   def __init__(self, initial_context: dict[str, Any] | None = None, handle: ParallelContextHandle | None = None):
     """
     Initializes the manager. If a handle is provided, it initializes in
-    proxy mode; otherwise, it starts a new server.
+    proxy mode; otherwise, it starts a new manager.
     """
     if handle:
       # --- PROXY MODE INITIALIZATION ---
-      # This instance is a client connecting to an existing server.
-      self._is_proxy = True
-      self._manager_server = None  # Proxies do not own the server process.
-
-      manager = _ParallelStateManager(address=handle.address)
-      manager.connect()
-      self._manager = manager
-
+      # This instance is a client wrapping objects from an existing server.
+      self._manager = None  # Proxies do not own the manager process.
+      self._shared_dict = handle._shared_dict
+      self._lock = handle._lock
     else:
-      # --- SERVER MODE INITIALIZATION ---
-      # This is the main instance that owns the server process.
-      self._is_proxy = False
-      manager = mp.Manager()  # type: ignore
-      _ParallelStateManager.register("get_dict", callable=lambda: manager.dict(initial_context or {}))
-      _ParallelStateManager.register("get_lock", callable=lambda: manager.Lock())
-
-      self._manager_server = _ParallelStateManager(address=("", 0))
-      self._manager_server.start()
-      self._manager = self._manager_server
-
-    # Common setup for both modes
-    self._shared_dict: DictProxy = self._manager.get_dict()  # type: ignore
-    self._lock: Lock = self._manager.get_lock()  # type: ignore
+      # --- MAIN MODE INITIALIZATION ---
+      # This instance owns the manager and its shared objects.
+      self._manager = mp.Manager()
+      self._shared_dict = self._manager.dict(initial_context or {})
+      self._lock = self._manager.Lock()
+
+    self._is_locked = False
+
+  def _lock_context(self) -> None:
+    """Acquire the lock for this context manager."""
+    if not self._is_locked:
+      self._lock.acquire()
+      self._is_locked = True
+
+  def _unlock_context(self) -> None:
+    """Release the lock for this context manager."""
+    if self._is_locked:
+      self._lock.release()
+      self._is_locked = False
+
+  def _execute_locked(self, operation: Callable[[], R]) -> R:
+    """A private helper to execute an operation within a lock."""
+    if not self._is_locked:
+      self._lock_context()
+      try:
+        return operation()
+      finally:
+        self._unlock_context()
+    else:
+      return operation()
 
   def get_handle(self) -> ParallelContextHandle:
     """
-    Returns a picklable handle for reconstruction in a worker.
-    Only the main server instance can generate handles.
+    Returns a picklable handle containing the shared dict and lock.
+    Only the main instance can generate handles.
     """
-    if self._is_proxy or not self._manager_server:
+    if not self._manager:
       raise TypeError("Cannot get a handle from a proxy context instance.")
 
-    return ParallelContextHandle(
-      address=self._manager_server.address,  # type: ignore
-      manager_class=self.__class__,  # Pass its own class for reconstruction
-    )
+    return ParallelContextHandle(self._shared_dict, self._lock)
 
   def shutdown(self) -> None:
     """
     Shuts down the background manager process.
-    This is a no-op for proxy instances, as only the main instance
-    should control the server's lifecycle.
+    This is a no-op for proxy instances.
     """
-    if not self._is_proxy and self._manager_server:
-      self._manager_server.shutdown()
+    if self._manager:
+      self._manager.shutdown()
 
   def __enter__(self) -> "ParallelContextManager":
     """Acquires the lock for use in a 'with' statement."""
-    self._lock.acquire()
+    self._lock_context()
     return self
 
   def __exit__(self, exc_type, exc_val, exc_tb) -> None:
     """Releases the lock."""
-    self._lock.release()
+    self._unlock_context()
 
   def __getitem__(self, key: str) -> Any:
-    with self._lock:
-      return self._shared_dict[key]
+    return self._execute_locked(lambda: self._shared_dict[key])
 
   def __setitem__(self, key: str, value: Any) -> None:
-    with self._lock:
-      self._shared_dict[key] = value
+    self._execute_locked(lambda: self._shared_dict.__setitem__(key, value))
 
   def __delitem__(self, key: str) -> None:
-    with self._lock:
-      del self._shared_dict[key]
+    self._execute_locked(lambda: self._shared_dict.__delitem__(key))
 
   def __iter__(self) -> Iterator[str]:
-    with self._lock:
-      return iter(list(self._shared_dict.keys()))
+    # Iteration needs to copy the keys to be safe across processes
+    return self._execute_locked(lambda: iter(list(self._shared_dict.keys())))
 
   def __len__(self) -> int:
-    with self._lock:
-      return len(self._shared_dict)
+    return self._execute_locked(lambda: len(self._shared_dict))

laygo/transformers/threaded.py

@@ -119,7 +119,7 @@ def __call__(self, data: Iterable[In], context: IContextManager | None = None) -
     Returns:
         An iterator over the transformed data.
     """
-    run_context = context if context is not None else self._default_context
+    run_context = context or self._default_context
 
     # Since threads share memory, we can pass the context manager directly.
     # No handle/proxy mechanism is needed, but the locking inside

tests/test_parallel_transformer.py

@@ -5,7 +5,8 @@
 
 from laygo import ErrorHandler
 from laygo import ParallelTransformer
-from laygo import PipelineContext
+from laygo.context import IContextManager
+from laygo.context import ParallelContextManager
 from laygo.transformers.parallel import createParallelTransformer
 from laygo.transformers.transformer import createTransformer
 
@@ -82,16 +83,18 @@ def test_tap_side_effects(self):
       assert sorted(side_effects) == [1, 2, 3, 4]
 
 
-def safe_increment(x: int, ctx: PipelineContext) -> int:
-  with ctx["lock"]:
+def safe_increment(x: int, ctx: IContextManager) -> int:
+  # Safe cast since we know ParallelContextManager implements context manager protocol
+  with ctx:  # type: ignore
     current_items = ctx["items"]
     time.sleep(0.001)
     ctx["items"] = current_items + 1
   return x * 2
 
 
-def update_stats(x: int, ctx: PipelineContext) -> int:
-  with ctx["lock"]:
+def update_stats(x: int, ctx: IContextManager) -> int:
+  # Safe cast since we know ParallelContextManager implements context manager protocol
+  with ctx:  # type: ignore
     ctx["total_sum"] += x
     ctx["item_count"] += 1
     ctx["max_value"] = max(ctx["max_value"], x)
@@ -103,25 +106,24 @@ class TestParallelTransformerContextSupport:
 
   def test_map_with_context(self):
     """Test map with context-aware function in concurrent execution."""
-    context = PipelineContext({"multiplier": 3})
+    context = ParallelContextManager({"multiplier": 3})
     transformer = createParallelTransformer(int).map(lambda x, ctx: x * ctx["multiplier"])
     result = list(transformer([1, 2, 3], context))
     assert result == [3, 6, 9]
 
-  def test_context_modification_with_locking(self):
-    """Test safe context modification with locking in concurrent execution."""
-    context = PipelineContext({"items": 0})
+  def test_context_aware_complex_operation(self):
+    """Test complex context-aware operations with shared state."""
+    context = ParallelContextManager({"multiplier": 3, "stats": {"total": 0, "count": 0}})
 
-    transformer = createParallelTransformer(int, max_workers=4, chunk_size=1).map(safe_increment)
-    data = list(range(1, 11))
+    transformer = createParallelTransformer(int, max_workers=2, chunk_size=2).map(update_stats)
+    data = [1, 2, 3, 4, 5]
     result = list(transformer(data, context))
 
-    assert sorted(result) == sorted([x * 2 for x in data])
-    assert context["items"] == len(data)
-
   def test_multiple_context_values_modification(self):
     """Test modifying multiple context values safely."""
-    context = PipelineContext({"total_sum": 0, "item_count": 0, "max_value": 0})
+    from laygo.context import ParallelContextManager
+
+    context = ParallelContextManager({"total_sum": 0, "item_count": 0, "max_value": 0})
 
     transformer = createParallelTransformer(int, max_workers=3, chunk_size=2).map(update_stats)
     data = [1, 5, 3, 8, 2, 7, 4, 6]