Repository: benchopt/benchopt
Pull request #722 (2024-06-03 09:39:47)

Saving final results of a solver
Hey there ! 

This PR add the possibility to save the final results of a benchmark to the filesystem. The data is saved as a `.pkl` file and a reference is added to the metadata of the benchmark. 

This is a much need capability for imaging problems, where getting visualisation of the results is complementary of quantitative metrics (PSNR, SSIM). 

It also opens the possibility to perform further processing on solver outputs (computing extra metrics, statistics, etc.)

### Checks before merging PR
- [x] added documentation for any new feature
- [x] added unit test
- [x] edited the [what's new](../../whatsnew.rst) (if applicable)


Patch:
diff --git a/.circleci/config.yml b/.circleci/config.yml
index 3cf46b0..1250bc3 100644
--- a/.circleci/config.yml
+++ b/.circleci/config.yml
@@ -135,7 +135,11 @@ jobs:
               git remote -v
               git fetch origin
               git reset --hard origin/main
-              git clean -xdf
+              mv .git ..
+              mv .gitignore ..
+              rm -rf *
+              mv ../.git .
+              mv ../.gitignore .
               echo "Deploying dev docs for ${CIRCLE_BRANCH}.";
               cp -a /tmp/build/html/* .;
               if [[ -z $(git diff) ]]; then
diff --git a/.github/workflows/expose_link.yml b/.github/workflows/expose_link.yml
index b353b5d..470bd0a 100644
--- a/.github/workflows/expose_link.yml
+++ b/.github/workflows/expose_link.yml
@@ -1,5 +1,10 @@
 name: Expose result artifact
 on: [status]
+# Cancel in-progress workflows when pushing
+# a new commit on the same branch
+concurrency:
+  group: ${{ github.workflow }}-${{ github.ref }}
+  cancel-in-progress: true
 jobs:
   circleci_result_artifact_redirector_job:
     runs-on: ubuntu-latest
diff --git a/.github/workflows/linting.yml b/.github/workflows/linting.yml
index 738cee4..58cf3ba 100644
--- a/.github/workflows/linting.yml
+++ b/.github/workflows/linting.yml
@@ -1,5 +1,4 @@
 name: linting
-
 on:
   push:
     branches:
@@ -10,6 +9,11 @@ on:
   create:
     tags:
       - '**'
+# Cancel in-progress workflows when pushing
+# a new commit on the same branch
+concurrency:
+  group: ${{ github.workflow }}-${{ github.ref }}
+  cancel-in-progress: true
 
 jobs:
   check-manifest:
diff --git a/.github/workflows/release.yml b/.github/workflows/release.yml
index 6329b6d..1d905c0 100644
--- a/.github/workflows/release.yml
+++ b/.github/workflows/release.yml
@@ -55,13 +55,14 @@ jobs:
 
     - name: Publish distribution 📦 to Test PyPI
       if: ${{ github.ref == 'refs/heads/main' && steps.wait-for-tests.outputs.conclusion == 'success' }}
-      uses: pypa/gh-action-pypi-publish@master
+      uses: pypa/gh-action-pypi-publish@release/v1
       with:
         password: ${{ secrets.TEST_PYPI_API_TOKEN }}
         repository_url: https://test.pypi.org/legacy/
 
     - name: Publish distribution 📦 to PyPI
       if: ${{ startsWith(github.ref, 'refs/tags') && steps.wait-for-tests.outputs.conclusion == 'success' }}
-      uses: pypa/gh-action-pypi-publish@master
+      uses: pypa/gh-action-pypi-publish@release/v1
       with:
         password: ${{ secrets.PYPI_API_TOKEN }}
+
diff --git a/benchopt/base.py b/benchopt/base.py
index 3b3923f..25a979b 100644
--- a/benchopt/base.py
+++ b/benchopt/base.py
@@ -360,6 +360,16 @@ class BaseObjective(ParametrizedNameMixin, DependenciesMixin, ABC):
       evaluated. This should be a dictionary where the keys correspond to the
       keyword arguments of `evaluate_result`.
 
+    Optionally, the `Solver` can implement the following methods to change its
+    behavior:
+
+    - `save_final_results(**result)`: Return the data to be saved from the
+       results of the solver. It will be saved as a `.pkl` file in the
+       `output/results` folder, and link to the benchmark results.
+
+    - `get_next(stop_val)`: Return the next iteration where the result will be
+      evaluated.
+
     This class is also used to specify information about the benchmark.
     In particular, it should have the following class attributes:
 
@@ -420,6 +430,20 @@ class BaseObjective(ParametrizedNameMixin, DependenciesMixin, ABC):
         """
         pass
 
+    def save_final_results(self, **solver_result):
+        """Save the final results of the solver.
+
+        Parameters
+        ----------
+        solver_result : dict
+            All values needed to compute the objective metrics. This dictionary
+            is retrieved by calling ``solver_result = Solver.get_result()``.
+        Returns
+        -------
+        dict of values to save
+        """
+        pass
+
     def __call__(self, result):
         """Used to call the evaluation of the objective.
 
diff --git a/benchopt/runner.py b/benchopt/runner.py
index 773a2a8..08608d4 100644
--- a/benchopt/runner.py
+++ b/benchopt/runner.py
@@ -1,9 +1,10 @@
 import time
 import inspect
+import pickle
 
 from datetime import datetime
 
-from joblib import Parallel, delayed
+from joblib import Parallel, delayed, hash
 
 from .callback import _Callback
 from .utils.sys_info import get_sys_info
@@ -93,6 +94,20 @@ def run_one_to_cvg(benchmark, objective, solver, meta, stopping_criterion,
     solver._warm_up()
 
     curve = []
+
+    # Augment the metadata with final_results if necessary.
+    base_method = getattr(
+        super(type(objective), objective),
+        'save_final_results', None
+    )
+
+    has_save_final_results = objective.save_final_results is not base_method
+    if has_save_final_results:
+        final_results = benchmark.get_output_folder() / 'final_results'
+        final_results /= f"{hash(meta)}.pkl"
+        final_results.parent.mkdir(exist_ok=True, parents=True)
+        meta["final_results"] = str(final_results)
+
     with exception_handler(output, pdb=pdb) as ctx:
 
         if solver._solver_strategy == "callback":
@@ -129,6 +144,12 @@ def run_one_to_cvg(benchmark, objective, solver, meta, stopping_criterion,
                 stop, ctx.status, stop_val = stopping_criterion.should_stop(
                     stop_val, curve
                 )
+        # Only run if save_final_results is defined in the objective.
+        if has_save_final_results:
+            to_save = objective.save_final_results(**solver.get_result())
+            if to_save is not None:
+                with open(meta["final_results"], 'wb') as f:
+                    pickle.dump(to_save, f)
 
     return curve, ctx.status
 
diff --git a/benchopt/stopping_criterion.py b/benchopt/stopping_criterion.py
index c321058..ff34a4c 100644
--- a/benchopt/stopping_criterion.py
+++ b/benchopt/stopping_criterion.py
@@ -212,7 +212,9 @@ class StoppingCriterion():
         n_eval = len(objective_list) - 1
         objective = objective_list[-1][self.key_to_monitor]
         delta_objective = self._prev_objective - objective
-        delta_objective /= abs(objective_list[0][self.key_to_monitor])
+        first_objective = objective_list[0][self.key_to_monitor]
+        if first_objective != 0:
+            delta_objective /= abs(first_objective)
         self._prev_objective = objective
 
         # default value for is_flat
@@ -430,7 +432,9 @@ class SufficientProgressCriterion(StoppingCriterion):
         # Compute the current objective and update best value
         objective = objective_list[-1][self.key_to_monitor]
         delta_objective = self._best_objective - objective
-        delta_objective /= abs(objective_list[0][self.key_to_monitor])
+        first_objective = objective_list[0][self.key_to_monitor]
+        if first_objective != 0:
+            delta_objective /= abs(first_objective)
         self._best_objective = min(
             objective, self._best_objective
         )
diff --git a/doc/names.inc b/doc/names.inc
index 493254c..f0c3399 100644
--- a/doc/names.inc
+++ b/doc/names.inc
@@ -25,3 +25,5 @@
 .. _Pierre Ablin: https://github.com/pierreablin
 
 .. _Christopher Marouani: https://github.com/chris-mrn
+
+.. _Pierre-Antoine Comby: https://github.com/paquiteau
diff --git a/doc/user_guide/advanced.rst b/doc/user_guide/advanced.rst
index e13a32a..3c0c46b 100644
--- a/doc/user_guide/advanced.rst
+++ b/doc/user_guide/advanced.rst
@@ -216,3 +216,10 @@ It assume that the python script is located at the same level as the benchmark f
 
 .. |SlurmExecutor| replace:: ``submitit.SlurmExecutor``
 .. _SlurmExecutor: https://github.com/facebookincubator/submitit/blob/main/submitit/slurm/slurm.py#L214
+
+.. _save_final_results:
+
+Saving Final Results of a Solver
+--------------------------------
+
+Using the `save_final_results(**results)` method of the objective function to retrieve the results to save. They are saved in `outputs/final_results/` directory and reference is added in the benchmark `.parquet` file.
diff --git a/doc/whats_new.rst b/doc/whats_new.rst
index 44e24d1..aac322c 100644
--- a/doc/whats_new.rst
+++ b/doc/whats_new.rst
@@ -13,6 +13,9 @@ Version 1.6 - In development
 API
 ~~~
 
+- Add a `save_final_results` method to Objective. If implemented it is run after the last solver iteration, to get desired outputs to be saved to file system.
+  By `Pierre-Antoine Comby`_ (:gh:`722`)
+
 - Add native way to do cross-validation in a benchmark with
   ``Objective.cv`` attribute that change split for each repetition.
   By `Christopher Marouani`_ and `Thomas Moreau`_ (:gh:`623`).


Test patch:
diff --git a/.github/workflows/test.yml b/.github/workflows/test.yml
index 946ce56..d99559c 100644
--- a/.github/workflows/test.yml
+++ b/.github/workflows/test.yml
@@ -1,6 +1,6 @@
 name: test
-
 on:
+
   push:
     branches:
       - main
@@ -14,6 +14,12 @@ on:
     # Run every day at 8:42am UTC.
     - cron:  '42 8 * * *'
 
+# Cancel in-progress workflows when pushing
+# a new commit on the same branch
+concurrency:
+  group: ${{ github.workflow }}-${{ github.ref }}
+  cancel-in-progress: true
+
 jobs:
   test_benchopt:
     name: Test
diff --git a/benchopt/tests/test_benchmark_features.py b/benchopt/tests/test_benchmark_features.py
index a9b125d..396fa93 100644
--- a/benchopt/tests/test_benchmark_features.py
+++ b/benchopt/tests/test_benchmark_features.py
@@ -113,6 +113,42 @@ def test_objective_no_cv(no_debug_test):
                 standalone_mode=False)
 
 
+def test_objective_save_final_results(no_debug_test):
+    save_final = """
+    from benchopt import BaseObjective
+
+    class Objective(BaseObjective):
+        name = "cross_val"
+
+        min_benchopt_version = "0.0.0"
+
+        def set_data(self, X, y): self.X, self.y = X, y
+        def get_one_result(self): return 0
+        def evaluate_result(self, beta): return dict(value=1)
+
+        def save_final_results(self, beta):
+            return "test_value"
+
+        def get_objective(self):
+            return dict(X=self.X, y=self.y, lmbd=1)
+
+    """
+
+    import pandas as pd
+    import pickle
+
+    with temp_benchmark(objective=save_final) as benchmark:
+        with CaptureRunOutput(delete_result_files=False) as out:
+            run([
+                str(benchmark.benchmark_dir),
+                *('-s python-pgd -d test-dataset -n 1 -r 1 --no-plot').split()
+            ],  standalone_mode=False)
+        data = pd.read_parquet(out.result_files[0])
+        with open(data.loc[0, "final_results"], "rb") as final_result_file:
+            final_results = pickle.load(final_result_file)
+    assert final_results == "test_value"
+
+
 def test_objective_cv_splitter(no_debug_test):
 
     objective = """from benchopt import BaseObjective, safe_import_context
diff --git a/benchopt/tests/test_stopping_criterion.py b/benchopt/tests/test_stopping_criterion.py
index b29ec01..bcc952d 100644
--- a/benchopt/tests/test_stopping_criterion.py
+++ b/benchopt/tests/test_stopping_criterion.py
@@ -294,3 +294,40 @@ def test_dual_strategy(no_debug_test):
                 run([str(benchmark.benchmark_dir),
                     *('-s test-solver -d test-dataset --no-plot').split()],
                     standalone_mode=False)
+
+
+def test_objective_equals_zero(no_debug_test):
+
+    objective = """from benchopt import BaseObjective
+
+        class Objective(BaseObjective):
+            name = "test_obj"
+            min_benchopt_version = "0.0.0"
+
+            def set_data(self, X, y): pass
+            def get_one_result(self): pass
+            def evaluate_result(self, beta): return dict(value=0)
+            def get_objective(self): return dict(X=0, y=0)
+    """
+
+    solver = """from benchopt import BaseSolver
+    from benchopt.stopping_criterion import SufficientDescentCriterion
+
+    class Solver(BaseSolver):
+        name = "test-solver"
+        stopping_criterion = SufficientDescentCriterion()
+        def set_objective(self, X, y): pass
+        def run(self, n_iter): pass
+        def get_result(self): return dict(beta=1)
+    """
+
+    with temp_benchmark(
+            objective=objective,
+            solvers=[solver]
+    ) as benchmark:
+        with CaptureRunOutput() as out:
+            run([str(benchmark.benchmark_dir),
+                *('-s test-solver -d test-dataset --no-plot -n 0').split()],
+                standalone_mode=False)
+
+    out.check_output('done', 1)