Stores per-suggestion latency in GP_UCB_PE suggestion timing metadata

PiperOrigin-RevId: 617964416

Author: vizier-team

vizier/_src/algorithms/designers/gp_ucb_pe.py

@@ -644,7 +644,7 @@ def suggest(
       jax.clear_caches()
 
     self._rng, rng = jax.random.split(self._rng, 2)
-    begin = datetime.datetime.now()
+    next_suggestion_start_time = datetime.datetime.now()
     data = self._trials_to_data(self._all_completed_trials)
     model = self._build_gp_model_and_optimize_parameters(data, rng)
     predictive = sp.UniformEnsemblePredictive(
@@ -806,13 +806,14 @@ def suggest(
           'params': f'{model.params}',
       })
       metadata.ns('timing').update(
-          {'time': f'{datetime.datetime.now() - begin}'}
+          {'time': f'{datetime.datetime.now() - next_suggestion_start_time}'}
       )
       suggestions.append(
           vz.TrialSuggestion(
               best_candidate.parameters, metadata=best_candidate.metadata
           )
       )
       active_trials.append(suggestions[-1].to_trial())
+      next_suggestion_start_time = datetime.datetime.now()
 
     return suggestions