really fix meta-data this time

Author: SimonBlanke

src/gradient_free_optimizers/distributed/_base.py

@@ -65,28 +65,31 @@ def __init__(self, n_workers: int):
         self.n_workers = n_workers
 
     @abstractmethod
-    def _distribute(self, func, params_batch: list[dict]) -> list[float]:
+    def _distribute(self, func, params_batch: list[dict]) -> list:
         """Evaluate func(params) for each params dict in the batch.
 
         This is the only method subclasses need to implement for
         synchronous (batch) evaluation. It receives the original
         (unwrapped) objective function and a list of parameter
-        dictionaries, and must return a list of scores in the same order.
+        dictionaries, and must return a list of results in the same order.
 
         The function follows GFO's convention: ``func(params_dict)`` where
         params_dict is a single dictionary, not keyword arguments.
 
         Parameters
         ----------
         func : callable
-            The original objective function with signature f(dict) -> float.
+            The original objective function with signature
+            f(dict) -> float or f(dict) -> (float, dict).
         params_batch : list[dict]
             Parameter dictionaries to evaluate.
 
         Returns
         -------
-        list[float]
-            Scores in the same order as params_batch.
+        list[float | tuple[float, dict]]
+            Results in the same order as params_batch. Each element is
+            either a plain score or a (score, metrics) tuple, matching
+            whatever the objective function returns.
         """
         ...
 
@@ -119,7 +122,7 @@ def _wait_any(self, futures) -> tuple:
 
         Only required for async backends (``_is_async = True``).
         Blocks until at least one future from the collection is ready,
-        then returns both the future object and its result value.
+        then returns both the future object and its raw result.
 
         Parameters
         ----------
@@ -128,8 +131,9 @@ def _wait_any(self, futures) -> tuple:
 
         Returns
         -------
-        tuple of (future, float)
-            The completed future object and its result (score).
+        tuple of (future, float | tuple[float, dict])
+            The completed future object and the objective function's
+            return value (a plain score or a (score, metrics) tuple).
             The future is returned so the caller can identify which
             submission completed (e.g., to look up the associated position).
         """
@@ -142,18 +146,20 @@ def distribute(self, func):
         """Decorator that wraps a single-point objective for batch evaluation.
 
         The decorated function accepts a list of parameter dicts and returns
-        a list of scores. It also carries metadata attributes that the
+        a list of results. It also carries metadata attributes that the
         optimizer's search loop uses to detect and configure batch mode.
 
         Parameters
         ----------
         func : callable
-            Objective function with signature f(params_dict) -> float.
+            Objective function with signature
+            f(params_dict) -> float or f(params_dict) -> (float, dict).
 
         Returns
         -------
         callable
-            Wrapped function with signature f(list[dict]) -> list[float].
+            Wrapped function with signature
+            f(list[dict]) -> list[float | tuple[float, dict]].
         """
 
         def wrapper(params_batch):

src/gradient_free_optimizers/search.py

@@ -136,13 +136,15 @@ def _iteration_batch(self, batch_size):
 
         # Check storage for cached results (preserving original order)
         scores = [None] * n_positions
+        metrics_list = [{}] * n_positions
         uncached_indices = []
 
         if self._storage is not None:
             for i, pos in enumerate(positions):
                 cached = self._storage.get(tuple(pos))
                 if cached is not None:
                     scores[i] = cached.score
+                    metrics_list[i] = cached.metrics
                 else:
                     uncached_indices.append(i)
         else:
@@ -156,16 +158,17 @@ def _iteration_batch(self, batch_size):
                 params_batch.append(self.conv.value2para(value))
 
             t_start = time.time()
-            new_scores = self._distributed_func(params_batch)
+            raw_results = self._distributed_func(params_batch)
             eval_time = time.time() - t_start
 
-            if self.optimum == "minimum":
-                new_scores = [-s for s in new_scores]
-
-            for idx, score in zip(uncached_indices, new_scores):
+            for idx, raw in zip(uncached_indices, raw_results):
+                score, metrics = self._unpack_result(raw)
+                if self.optimum == "minimum":
+                    score = -score
                 scores[idx] = score
+                metrics_list[idx] = metrics
                 if self._storage is not None:
-                    self._storage.put(tuple(positions[idx]), Result(score, {}))
+                    self._storage.put(tuple(positions[idx]), Result(score, metrics))
 
             per_eval_time = eval_time / len(uncached_indices)
         else:
@@ -179,25 +182,50 @@ def _iteration_batch(self, batch_size):
         uncached_set = set(uncached_indices)
         for i, (pos, score) in enumerate(zip(positions, scores)):
             et = per_eval_time if i in uncached_set else 0
-            self._track_evaluation(pos, score, et, per_iter_time)
+            self._track_evaluation(
+                pos, score, et, per_iter_time, metrics=metrics_list[i]
+            )
+
+    @staticmethod
+    def _unpack_result(raw):
+        """Separate a worker return value into score and metrics.
+
+        Objective functions may return a plain float or a (float, dict) tuple.
+        The backends pass through whatever the function returns, so we unpack
+        here at the boundary between worker results and the search loop.
+        """
+        if isinstance(raw, tuple):
+            return raw[0], raw[1]
+        return raw, {}
 
-    def _track_evaluation(self, pos, score, eval_time=0, iter_time=0):
+    def _track_evaluation(self, pos, score, eval_time=0, iter_time=0, metrics=None):
         """Record a single evaluation result across all tracking systems."""
+        if metrics is None:
+            metrics = {}
         self.eval_times.append(eval_time)
         self.iter_times.append(iter_time)
-        self.results_manager.add(Result(score, {}), pos)
+        self.results_manager.add(Result(score, metrics), pos)
         self.pos_l.append(pos)
         self.score_l.append(score)
         self.p_bar.update(score, pos, self.nth_iter)
-        self._tracker.track(pos, score, {}, is_init=False)
+        self._last_metrics = metrics
+        self._tracker.track(pos, score, metrics, is_init=False)
+        # Feed each evaluation to the stopper so early_stopping counts
+        # individual evaluations, not batches
+        self.stopper.update(score, self.p_bar.score_best, self._iter)
         self.n_iter_total += 1
         self.n_iter_search += 1
         self._iter += 1
 
     def _check_stop(self, n_evaluated):
         """Check stopping conditions and run callbacks. Returns True if should stop."""
-        current_score = self.score_l[-1] if self.score_l else -math.inf
-        self.stopper.update(current_score, self.p_bar.score_best, n_evaluated - 1)
+        # In serial mode, the stopper needs per-evaluation updates here.
+        # In distributed mode, _track_evaluation already updated the stopper
+        # for each evaluation in the batch, giving correct granularity for
+        # early_stopping counters.
+        if not self._is_distributed:
+            current_score = self.score_l[-1] if self.score_l else -math.inf
+            self.stopper.update(current_score, self.p_bar.score_best, n_evaluated - 1)
 
         if self.stopper.should_stop():
             if "debug_stop" in self.verbosity:
@@ -248,7 +276,7 @@ def _submit_one():
                 cached = self._storage.get(tuple(pos))
                 if cached is not None:
                     self._evaluate_batch([pos], [cached.score])
-                    self._track_evaluation(pos, cached.score)
+                    self._track_evaluation(pos, cached.score, metrics=cached.metrics)
                     n_evaluated += 1
                     return 1
 
@@ -267,19 +295,20 @@ def _submit_one():
         # Process results as they arrive
         while pending and n_evaluated < n_iter:
             t_iter = time.time()
-            completed, score = backend._wait_any(list(pending.keys()))
+            completed, raw = backend._wait_any(list(pending.keys()))
             pos = pending.pop(completed)
 
+            score, metrics = self._unpack_result(raw)
             if self.optimum == "minimum":
                 score = -score
 
             if self._storage is not None:
-                self._storage.put(tuple(pos), Result(score, {}))
+                self._storage.put(tuple(pos), Result(score, metrics))
 
             self.nth_iter = n_evaluated
             self._evaluate_batch([pos], [score])
             iter_time = time.time() - t_iter
-            self._track_evaluation(pos, score, iter_time, iter_time)
+            self._track_evaluation(pos, score, iter_time, iter_time, metrics=metrics)
             n_evaluated += 1
 
             if self._check_stop(n_evaluated):
@@ -315,13 +344,15 @@ def _run_batch_async(self, n_iter, nth_trial):
 
             # Check storage cache (same logic as sync _iteration_batch)
             scores = [None] * n_positions
+            metrics_list = [{}] * n_positions
             uncached_indices = []
 
             if self._storage is not None:
                 for i, pos in enumerate(positions):
                     cached = self._storage.get(tuple(pos))
                     if cached is not None:
                         scores[i] = cached.score
+                        metrics_list[i] = cached.metrics
                     else:
                         uncached_indices.append(i)
             else:
@@ -339,15 +370,17 @@ def _run_batch_async(self, n_iter, nth_trial):
                 # Collect all results (async within batch)
                 t_start = time.time()
                 while futures:
-                    completed, score = backend._wait_any(list(futures.keys()))
+                    completed, raw = backend._wait_any(list(futures.keys()))
                     idx = futures.pop(completed)
 
+                    score, metrics = self._unpack_result(raw)
                     if self.optimum == "minimum":
                         score = -score
 
                     scores[idx] = score
+                    metrics_list[idx] = metrics
                     if self._storage is not None:
-                        self._storage.put(tuple(positions[idx]), Result(score, {}))
+                        self._storage.put(tuple(positions[idx]), Result(score, metrics))
 
                 per_eval_time = (time.time() - t_start) / len(uncached_indices)
             else:
@@ -361,7 +394,9 @@ def _run_batch_async(self, n_iter, nth_trial):
             uncached_set = set(uncached_indices)
             for i, (pos, score) in enumerate(zip(positions, scores)):
                 et = per_eval_time if i in uncached_set else 0
-                self._track_evaluation(pos, score, et, per_iter_time)
+                self._track_evaluation(
+                    pos, score, et, per_iter_time, metrics=metrics_list[i]
+                )
 
             n_evaluated += n_positions
 
@@ -681,27 +716,9 @@ def _init_search(
             # Extract original function for single-point use during init
             objective_function = objective_function._gfo_original_func
 
-            # The objective may return (score, metrics) tuples, but the
-            # distributed batch interface only passes scores between
-            # workers and the search loop. Normalize here so all
-            # distributed paths (sync batch, true-async, batch-async)
-            # receive plain floats.
-            _raw_distributed = self._original_func
-
-            def _normalize_return(params):
-                out = _raw_distributed(params)
-                if isinstance(out, tuple):
-                    return out[0]
-                return out
-
-            self._original_func = _normalize_return
-
             if catch:
                 self._original_func = wrap_with_catch(self._original_func, catch)
-
-            # Rebuild the wrapper so the sync batch path also uses the
-            # normalized (and optionally catch-wrapped) function
-            self._distributed_func = self._backend.distribute(self._original_func)
+                self._distributed_func = self._backend.distribute(self._original_func)
         else:
             self._batch_size = None
             self._backend = None