Repository: benchopt/benchopt
Pull request #713 (2024-05-16 20:58:25)

FIX benchmark_utils pickling
Follow up of #712 which was merged a bit quickly.
This should fix the issue in scikit-adaptation/skada-bench#60 with missing `benchmark_utils` module.

tomMoral: Let's just wait for confirmation that this fixes the issue of @antoinecollas  :)
Otherwise you can merge ;)
antoinecollas: It works! Thanks a lot Thomas 🙏🏻
tomMoral: @agramfort feel free to merge once it is green :wink: 
codecov-commenter: ## [Codecov](https://app.codecov.io/gh/benchopt/benchopt/pull/713?dropdown=coverage&src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=benchopt) Report
Attention: Patch coverage is `20.00000%` with `8 lines` in your changes are missing coverage. Please review.
> Project coverage is 49.42%. Comparing base [(`9f5c9f7`)](https://app.codecov.io/gh/benchopt/benchopt/commit/9f5c9f7a1f23b73da49441465bd5d5345d000968?dropdown=coverage&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=benchopt) to head [(`b94a043`)](https://app.codecov.io/gh/benchopt/benchopt/pull/713?dropdown=coverage&src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=benchopt).

> :exclamation: **Current head b94a043 differs from pull request most recent head 322f586**
> 
> Please [upload](https://docs.codecov.com/docs/codecov-uploader) reports for the commit 322f586 to get more accurate results.

:exclamation: Your organization needs to install the [Codecov GitHub app](https://github.com/apps/codecov/installations/select_target) to enable full functionality.

<details><summary>Additional details and impacted files</summary>


```diff
@@            Coverage Diff             @@
##             main     #713      +/-   ##
==========================================
- Coverage   49.51%   49.42%   -0.10%     
==========================================
  Files          45       45              
  Lines        3009     3019      +10     
  Branches      570      572       +2     
==========================================
+ Hits         1490     1492       +2     
- Misses       1376     1383       +7     
- Partials      143      144       +1     
```



</details>

Patch:
diff --git a/benchopt/utils/safe_import.py b/benchopt/utils/safe_import.py
index 814fd66..19d1bf7 100644
--- a/benchopt/utils/safe_import.py
+++ b/benchopt/utils/safe_import.py
@@ -3,6 +3,8 @@ import warnings
 import importlib
 from pathlib import Path
 
+from joblib.externals import cloudpickle
+
 from ..config import RAISE_INSTALL_ERROR
 
 SKIP_IMPORT = False
@@ -38,6 +40,7 @@ def set_benchmark_module(benchmark_dir):
         module = importlib.util.module_from_spec(spec)
         sys.modules[PACKAGE_NAME] = module
         spec.loader.exec_module(module)
+        cloudpickle.register_pickle_by_value(module)
     elif module_file.parent.exists():
         warnings.warn(
             "Folder `benchmark_utils` exists but is missing `__init__.py`. "
diff --git a/benchopt/utils/temp_benchmark.py b/benchopt/utils/temp_benchmark.py
index 5c1fa04..75964e8 100644
--- a/benchopt/utils/temp_benchmark.py
+++ b/benchopt/utils/temp_benchmark.py
@@ -26,7 +26,8 @@ dummy_datasets = [
 
 @contextlib.contextmanager
 def temp_benchmark(
-        objective=None, datasets=None, solvers=None, config=None
+        objective=None, datasets=None, solvers=None, config=None,
+        benchmark_utils=None
 ):
     """Create Benchmark in a temporary folder, for test purposes.
 
@@ -44,6 +45,7 @@ def temp_benchmark(
     config: str | None (default=None)
         Content of configuration file for running the Benchmark. If None,
         no config file is created.
+    benchmark_utils: dict(fname->str) | None (default=None)
     """
     if objective is None:
         objective = dummy_objective
@@ -75,4 +77,13 @@ def temp_benchmark(
             with open(temp_path / "config.yml", "w") as f:
                 f.write(config)
 
+        if benchmark_utils is not None:
+            benchmark_utils_dir = (temp_path / "benchmark_utils")
+            benchmark_utils_dir.mkdir()
+            (benchmark_utils_dir / "__init__.py").touch()
+            for fname, content in benchmark_utils.items():
+                fname = (benchmark_utils_dir / fname).with_suffix(".py")
+                with open(fname, "w") as f:
+                    f.write(inspect.cleandoc(content))
+
         yield Benchmark(temp_path)
diff --git a/doc/whats_new.rst b/doc/whats_new.rst
index 187b798..4867902 100644
--- a/doc/whats_new.rst
+++ b/doc/whats_new.rst
@@ -27,6 +27,11 @@ CLI
 - Add ``--collect`` option to allow gathering results which are already
   in cache in a single parquet file. By `Thomas Moreau`_ (:gh:`710`)
 
+FIX
+~~~
+
+- Fix pickling of dynamic modules to allow for nested parallelism in
+  distributed runs. By `Thomas Moreau`_ (:gh:`713`)
 
 .. _changes_1_5_1:
 


Test patch:
diff --git a/benchopt/utils/tests/test_dynamic_module.py b/benchopt/utils/tests/test_dynamic_module.py
new file mode 100644
index 0000000..f43c8ce
--- /dev/null
+++ b/benchopt/utils/tests/test_dynamic_module.py
@@ -0,0 +1,42 @@
+from joblib import Parallel, delayed
+
+from benchopt.tests.utils import CaptureRunOutput
+from benchopt.utils.temp_benchmark import temp_benchmark
+
+
+def test_pickling_dynamic_module():
+    # Make sure the dynamic modules can be pickled by joblib. In particular,
+    # this is necessary for nested parallelism in distributed context.
+    #
+    # This test check that the module containing a solver can be pickled, and
+    # that the dynamic benchmark_utils module can be retrieved in the process.
+
+    solver = """
+    from benchopt import BaseSolver
+    from benchmark_utils.test1 import func1
+
+    class Solver(BaseSolver):
+        name = "test-solver"
+        def set_objective(self, X, y, reg): self.X, self.y = X, y
+        def run(self, _): func1()
+        def get_result(self): return dict(beta=1)
+    """
+
+    with temp_benchmark(
+            solvers=[solver],
+            benchmark_utils={'test1': 'def func1(): print("FUNC1")'}
+    ) as benchmark:
+        # Get the solver from the list of tuples
+        Solver, _ = benchmark.check_solver_patterns(["test-solver"])[0]
+        assert Solver.is_installed()
+
+        with CaptureRunOutput() as out:
+            Solver.run(None, None)
+        out.check_output("FUNC1", repetition=1)
+
+        # This will fail if the benchmark_utils module is not pickled by value
+        # by cloudpickle.
+        with CaptureRunOutput():
+            Parallel(n_jobs=2)(
+                delayed(Solver.run)(None, None) for _ in range(2)
+            )