fix(probes,generators): raise GarakException on unpicklable parallel worker payload

garak/generators/base.py

@@ -4,6 +4,7 @@
 """
 
 import logging
+import pickle
 import random
 import re
 from typing import List, Union
@@ -210,6 +211,13 @@ def generate(
                         raise GarakException(msg) from o
                     else:
                         raise (o)
+                except pickle.PicklingError as p:
+                    msg = (
+                        "A generator attribute could not be sent to a parallel worker process. "
+                        "Try setting parallel_requests to 1."
+                    )
+                    logging.critical(msg)
+                    raise GarakException(msg) from p
                 finally:
                     if pool is not None:
                         pool.close()

garak/probes/base.py

@@ -12,6 +12,7 @@
 import copy
 import json
 import logging
+import pickle
 from collections.abc import Iterable
 import random
 from typing import Iterable, List, Set, Union
@@ -363,6 +364,13 @@ def _execute_all(self, attempts) -> Iterable[garak.attempt.Attempt]:
                     raise GarakException(msg) from o
                 else:
                     raise (o)
+            except pickle.PicklingError as p:
+                msg = (
+                    "A probe or generator attribute could not be sent to a parallel worker "
+                    "process. Try setting parallel_attempts to 1."
+                )
+                logging.critical(msg)
+                raise GarakException(msg) from p
             finally:
                 if attempt_pool is not None:
                     attempt_pool.close()

tests/generators/test_generators.py

@@ -9,6 +9,7 @@
 
 from garak import _plugins
 from garak import _config
+from garak.exception import GarakException
 
 from garak.attempt import Message, Turn, Conversation
 from garak.generators.base import Generator
@@ -40,6 +41,32 @@ def test_parallel_requests():
     ), "All generated Message texts should be non-empty"
 
 
+def _unpicklable_stub():
+    """Module-level stub whose __module__ is corrupted per-test to force a
+    real pickle.PicklingError (import-of-module-failed), matching the
+    dynamically-loaded-plugin failure mode reported in #361."""
+
+
+def test_parallel_requests_unpicklable_attribute_raises_garak_exception():
+    """A generator holding a non-picklable attribute must fail with a clear
+    GarakException rather than letting Pool._handle_tasks crash with a raw
+    _pickle.PicklingError (#361)."""
+    parallel_count = 2
+    _config.system.parallel_requests = parallel_count
+    _config.system.max_workers = parallel_count
+
+    g = _plugins.load_plugin("generators.test.Lipsum")
+    original_module = _unpicklable_stub.__module__
+    _unpicklable_stub.__module__ = "nonexistent_module_for_pickling_test"
+    g.unpicklable = _unpicklable_stub
+    try:
+        prompt = Conversation(Turn("user", [Message("this is a test")]))
+        with pytest.raises(GarakException):
+            g.generate(prompt=prompt, generations_this_call=3)
+    finally:
+        _unpicklable_stub.__module__ = original_module
+
+
 @pytest.mark.parametrize("classname", GENERATORS)
 def test_generator_structure(classname):
 

tests/test_internal_structures.py

@@ -16,6 +16,7 @@
 import garak.attempt
 import garak.buffs.base
 import garak.evaluators.base
+import garak.exception
 import garak.harnesses.base
 
 from garak.detectors.mitigation import MitigationBypass
@@ -66,6 +67,36 @@ def test_generator_consume_attempt_generator():
     ), "there should be the same number of attempts in the passed generator as results returned in _execute_all"
 
 
+def _unpicklable_stub():
+    """Module-level stub whose __module__ is corrupted per-test to force a
+    real pickle.PicklingError (import-of-module-failed), matching the
+    dynamically-loaded-plugin failure mode reported in #361."""
+
+
+def test_execute_all_unpicklable_probe_raises_garak_exception():
+    """A probe holding a non-picklable attribute must fail with a clear
+    GarakException rather than letting Pool._handle_tasks crash with a raw
+    _pickle.PicklingError (#361)."""
+    garak._config.system.parallel_attempts = 2
+    garak._config.system.max_workers = 2
+
+    p = garak._plugins.load_plugin("probes.test.Blank")
+    g = garak._plugins.load_plugin("generators.test.Blank")
+    p.generator = g
+    original_module = _unpicklable_stub.__module__
+    _unpicklable_stub.__module__ = "nonexistent_module_for_pickling_test"
+    p.unpicklable = _unpicklable_stub
+    attempts = [
+        garak.attempt.Attempt(prompt=garak.attempt.Message(text=str(i), lang="*"))
+        for i in range(2)
+    ]
+    try:
+        with pytest.raises(garak.exception.GarakException):
+            p._execute_all(attempts)
+    finally:
+        _unpicklable_stub.__module__ = original_module
+
+
 def test_attempt_outputs_can_consume_generator():
     a = garak.attempt.Attempt(prompt=garak.attempt.Message(text="fish", lang="*"))
     count = 5