Internal clean up.

PiperOrigin-RevId: 614747664

vizier/_src/algorithms/designers/eagle_strategy/eagle_strategy_utils_test.py

@@ -281,7 +281,7 @@ def test_replace_trial_metric_name(self):
  trial.complete(measurement=vz.Measurement(metrics={'obj123': 1123.3}))
  new_trial = utils.standardize_trial_metric_name(trial)
  self.assertEqual(
- new_trial.final_measurement.metrics['objective'].value, 1123.3
+ new_trial.final_measurement_or_die.metrics['objective'].value, 1123.3
  )
  self.assertEqual(new_trial.parameters['f1'].value, 0.0)
  self.assertEqual(new_trial.metadata.ns('eagle')['parent_fly_id'], '123')

vizier/_src/algorithms/designers/eagle_strategy/serialization_test.py

@@ -60,8 +60,10 @@ def test_restore_pool(self, x_values, obj_values):
  firefly.trial._infeasibility_reason,
  )
  self.assertEqual(
- restored_firefly.trial.final_measurement.metrics['objective'].value,
- firefly.trial.final_measurement.metrics['objective'].value,
+ restored_firefly.trial.final_measurement_or_die.metrics[
+ 'objective'
+ ].value,
+ firefly.trial.final_measurement_or_die.metrics['objective'].value,
  )
 
  def test_restore_rng(self):

vizier/_src/algorithms/designers/scalarizing_designer.py

@@ -77,13 +77,13 @@ def update(
  ) -> None:
  for trial in completed.trials:
  objectives = [
- trial.final_measurement.metrics.get_value(
+ trial.final_measurement_or_die.metrics.get_value(
  config.name, default=jnp.nan
  )
  for config in self._objectives
  ]
  # Simply append the scalarized value.
- trial.final_measurement.metrics[self._scalarized_metric_name] = (
+ trial.final_measurement_or_die.metrics[self._scalarized_metric_name] = (
  self._scalarizer(jnp.array(objectives))
  )
 

vizier/_src/algorithms/designers/scalarizing_designer_test.py

@@ -135,7 +135,7 @@ def eagle_designer_factory(ps, seed):
  all_active=vza.ActiveTrials(trials=[]),
  )
  self.assertTrue(
- jnp.isnan(trial.final_measurement.metrics['scalarized'].value)
+ jnp.isnan(trial.final_measurement_or_die.metrics['scalarized'].value)
  )
 
 

vizier/_src/algorithms/optimizers/vectorized_base_test.py

@@ -224,7 +224,9 @@ def test_best_candidates_count_is_1(self, use_fori):
  self.assertEqual(best_candidates[0].parameters['f1'].value, 0.5)
  self.assertEqual(best_candidates[0].parameters['f2'].value, 0.5)
  self.assertAlmostEqual(
- best_candidates[0].final_measurement.metrics['acquisition'].value,
+ best_candidates[0]
+ .final_measurement_or_die.metrics['acquisition']
+ .value,
  -((0.5 - 0.52) ** 2),
  )
 
@@ -251,21 +253,27 @@ def test_best_candidates_count_is_3(self, use_fori):
  self.assertAlmostEqual(best_candidates[0].parameters['f1'].value, 0.5)
  self.assertAlmostEqual(best_candidates[0].parameters['f2'].value, 0.5)
  self.assertAlmostEqual(
- best_candidates[0].final_measurement.metrics['acquisition'].value,
+ best_candidates[0]
+ .final_measurement_or_die.metrics['acquisition']
+ .value,
  -((0.5 - 0.52) ** 2),
  )
  # check 2nd best candidate
  self.assertAlmostEqual(best_candidates[1].parameters['f1'].value, 0.6)
  self.assertAlmostEqual(best_candidates[1].parameters['f2'].value, 0.6)
  self.assertAlmostEqual(
- best_candidates[1].final_measurement.metrics['acquisition'].value,
+ best_candidates[1]
+ .final_measurement_or_die.metrics['acquisition']
+ .value,
  -((0.6 - 0.52) ** 2),
  )
  # check 3rd best candidate
  self.assertAlmostEqual(best_candidates[2].parameters['f1'].value, 0.4)
  self.assertAlmostEqual(best_candidates[2].parameters['f2'].value, 0.4)
  self.assertAlmostEqual(
- best_candidates[2].final_measurement.metrics['acquisition'].value,
+ best_candidates[2]
+ .final_measurement_or_die.metrics['acquisition']
+ .value,
  -((0.4 - 0.52) ** 2),
  )
 

vizier/_src/algorithms/testing/comparator_runner.py

@@ -170,13 +170,15 @@ def assert_optimizer_better_simple_regret(
  res = baseline_optimizer(score_fn, count=1, seed=random.PRNGKey(i)) # pytype: disable=wrong-arg-types
  trial = vb.best_candidates_to_trials(res, converter)
  baseline_obj_values.append(
- trial[0].final_measurement.metrics['acquisition'].value)
+ trial[0].final_measurement_or_die.metrics['acquisition'].value
+ )
 
  for i in range(self.candidate_num_repeats):
  res = candidate_optimizer(score_fn, count=1, seed=random.PRNGKey(i)) # pytype: disable=wrong-arg-types
  trial = vb.best_candidates_to_trials(res, converter)
  candidate_obj_values.append(
- trial[0].final_measurement.metrics['acquisition'].value)
+ trial[0].final_measurement_or_die.metrics['acquisition'].value
+ )
 
  self._conclude_test(baseline_obj_values, candidate_obj_values)
 
@@ -199,7 +201,7 @@ def _run_one(benchmark_state_factory: benchmarks.BenchmarkStateFactory,
  best_trial = benchmark_state.algorithm.supporter.GetBestTrials(count=1)[0]
  metric_name = benchmark_state.experimenter.problem_statement(
  ).single_objective_metric_name
- return best_trial.final_measurement.metrics[metric_name].value
+ return best_trial.final_measurement_or_die.metrics[metric_name].value
 
  baseline_obj_values = []
  candidate_obj_values = []

vizier/_src/benchmarks/experimenters/combo_experimenter_test.py

@@ -58,7 +58,7 @@ def test_experimenters(self, experimenter_class, objective_min,
  logging.info('Evaluated Trial: %s', trial)
  self.assertEqual(trial.status, pyvizier.TrialStatus.COMPLETED)
  metric_name = problem_statement.metric_information.item().name
- eval_objective = trial.final_measurement.metrics[metric_name].value
+ eval_objective = trial.final_measurement_or_die.metrics[metric_name].value
  self.assertLessEqual(eval_objective, objective_max)
  self.assertGreaterEqual(eval_objective, objective_min)
 

vizier/_src/benchmarks/experimenters/discretizing_experimenter_test.py

@@ -76,7 +76,8 @@ def testNumpyExperimenter(self, func):
  metric_name = exptr.problem_statement().metric_information.item().name
  self.assertAlmostEqual(
  func(np.array([0.0, 1.0, 1.5])),
- t.final_measurement.metrics[metric_name].value)
+ t.final_measurement_or_die.metrics[metric_name].value,
+ )
  self.assertEqual(t.status, pyvizier.TrialStatus.COMPLETED)
  self.assertDictEqual(t.parameters.as_dict(), parameters)
 

vizier/_src/benchmarks/experimenters/hpob_experimenter_test.py

@@ -61,18 +61,25 @@ def test_all(self, **kwargs):
 
  experimenter.evaluate([problem_and_trials.trials[0]])
 
- logging.info('First five trial metrics: %s', [
- t.final_measurement.metrics[hpob_experimenter.METRIC_NAME].value
- for t in problem_and_trials.trials[:5]
- ])
+ logging.info(
+ 'First five trial metrics: %s',
+ [
+ t.final_measurement_or_die.metrics[
+ hpob_experimenter.METRIC_NAME
+ ].value
+ for t in problem_and_trials.trials[:5]
+ ],
+ )
  logging.info('objective: %f', objective)
  logging.info('evaluated: %s', problem_and_trials.trials[0])
 
  self.assertAlmostEqual(
- problem_and_trials.trials[0].final_measurement.metrics[
- hpob_experimenter.METRIC_NAME].value,
+ problem_and_trials.trials[0]
+ .final_measurement_or_die.metrics[hpob_experimenter.METRIC_NAME]
+ .value,
  objective,
- places=1)
+ places=1,
+ )
 
  break
 

vizier/_src/benchmarks/experimenters/l1_categorical_experimenter_test.py

@@ -32,15 +32,19 @@ def test_evaluate_optimum(self):
  num_categories=[2, 2], optimum=optimum)
  suggestion = vz.Trial(parameters={'c0': '0', 'c1': '1'})
  exptr.evaluate([suggestion])
- self.assertEqual(suggestion.final_measurement.metrics['objective'].value, 0)
+ self.assertEqual(
+ suggestion.final_measurement_or_die.metrics['objective'].value, 0
+ )
 
  def test_evaluate_non_optimum(self):
  optimum = [0, 1]
  exptr = l1_categorical_experimenter.L1CategorialExperimenter(
  num_categories=[2, 2], optimum=optimum)
  suggestion = vz.Trial(parameters={'c0': '1', 'c1': '0'})
  exptr.evaluate([suggestion])
- self.assertEqual(suggestion.final_measurement.metrics['objective'].value, 2)
+ self.assertEqual(
+ suggestion.final_measurement_or_die.metrics['objective'].value, 2
+ )
 
  @parameterized.parameters({'num_categories': [10, 3, 2]},
  {'num_categories': [10, 2, 10]},

vizier/_src/benchmarks/experimenters/multiobjective_experimenter_test.py

@@ -50,11 +50,11 @@ def test_mulitobjective_numpy(self):
  exptr.evaluate([t])
  self.assertAlmostEqual(
  func1(np.array([0.0, 1.0])),
- t.final_measurement.metrics['m1'].value,
+ t.final_measurement_or_die.metrics['m1'].value,
  )
  self.assertAlmostEqual(
  func2(np.array([0.0, 1.0])),
- t.final_measurement.metrics['m2'].value,
+ t.final_measurement_or_die.metrics['m2'].value,
  )
  self.assertEqual(t.status, pyvizier.TrialStatus.COMPLETED)
 

vizier/_src/benchmarks/experimenters/nasbench101_experimenter_test.py

@@ -42,7 +42,9 @@ def test_experimenter(self):
  self.assertEqual(trial.status, pyvizier.TrialStatus.COMPLETED)
  if not trial.infeasible:
  metric_name = problem_statement.metric_information.item().name
- eval_objective = trial.final_measurement.metrics[metric_name].value
+ eval_objective = trial.final_measurement_or_die.metrics[
+ metric_name
+ ].value
  self.assertGreaterEqual(eval_objective, 0.0)
  self.assertLessEqual(eval_objective, 100.0)
 

vizier/_src/benchmarks/experimenters/nasbench201_experimenter_test.py

@@ -41,7 +41,7 @@ def test_experimenter(self):
  logging.info('Evaluated Trial: %s', trial)
  self.assertEqual(trial.status, pyvizier.TrialStatus.COMPLETED)
  metric_name = problem_statement.metric_information.item().name
- eval_objective = trial.final_measurement.metrics[metric_name].value
+ eval_objective = trial.final_measurement_or_die.metrics[metric_name].value
  self.assertGreaterEqual(eval_objective, 0.0)
  self.assertLessEqual(eval_objective, 100.0)
 

vizier/_src/benchmarks/experimenters/noisy_experimenter_test.py

@@ -59,14 +59,15 @@ def testDeterministicNoiseApply(self, func):
 
  exptr.evaluate([t])
  metric_name = exptr.problem_statement().metric_information.item().name
- unnoised_value = t.final_measurement.metrics[metric_name].value
+ unnoised_value = t.final_measurement_or_die.metrics[metric_name].value
 
  noisy_exptr.evaluate([t])
- noised_value = t.final_measurement.metrics[metric_name].value
+ noised_value = t.final_measurement_or_die.metrics[metric_name].value
  self.assertEqual(unnoised_value - 1, noised_value)
  self.assertEqual(
  unnoised_value,
- t.final_measurement.metrics[metric_name + '_before_noise'].value)
+ t.final_measurement_or_die.metrics[metric_name + '_before_noise'].value,
+ )
 
  @parameterized.named_parameters(
  ('NO_NOISE', 'NO_NOISE', 1e-5),
@@ -95,13 +96,13 @@ def testGaussianNoiseApply(self, noise: str, delta: float):
 
  exptr.evaluate([t])
  metric_name = exptr.problem_statement().metric_information.item().name
- unnoised_value = t.final_measurement.metrics[metric_name].value
+ unnoised_value = t.final_measurement_or_die.metrics[metric_name].value
 
  noisy_exptr.evaluate([t])
- noised_value1 = t.final_measurement.metrics[metric_name].value
+ noised_value1 = t.final_measurement_or_die.metrics[metric_name].value
 
  noisy_exptr.evaluate([t])
- noised_value2 = t.final_measurement.metrics[metric_name].value
+ noised_value2 = t.final_measurement_or_die.metrics[metric_name].value
 
  # Seldom noise is only injected sporadically.
  if 'SELDOM' not in noise and noise != 'NO_NOISE':
@@ -131,7 +132,7 @@ def testSeedDeterminism(self):
  for _ in range(10):
  noisy_exptr.evaluate([t])
  noise_value_sequence.append(
- t.final_measurement.metrics[metric_name].value
+ t.final_measurement_or_die.metrics[metric_name].value
  )
 
  # Global NP seed should not affect randomness.
@@ -144,7 +145,7 @@ def testSeedDeterminism(self):
  for _ in range(10):
  noisy_exptr.evaluate([t])
  noise_value_sequence_after.append(
- t.final_measurement.metrics[metric_name].value
+ t.final_measurement_or_die.metrics[metric_name].value
  )
  self.assertSequenceAlmostEqual(
  noise_value_sequence, noise_value_sequence_after

vizier/_src/benchmarks/experimenters/normalizing_experimenter_test.py

@@ -68,7 +68,7 @@ def testNormalizationApply(self, func):
  metric_name = exptr.problem_statement().metric_information.item().name
 
  normalizing_exptr.evaluate([t])
- normalized_value = t.final_measurement.metrics[metric_name].value
+ normalized_value = t.final_measurement_or_die.metrics[metric_name].value
  self.assertBetween(normalized_value, -10, 10)
 
 

vizier/_src/benchmarks/experimenters/numpy_experimenter_test.py

@@ -57,7 +57,8 @@ def testNumpyExperimenter(self, func):
  metric_name = exptr.problem_statement().metric_information.item().name
  self.assertAlmostEqual(
  func(np.array([0.0, 1.0])),
- t.final_measurement.metrics[metric_name].value)
+ t.final_measurement_or_die.metrics[metric_name].value,
+ )
  self.assertEqual(t.status, pyvizier.TrialStatus.COMPLETED)
 
  def testNonFinite(self):
@@ -79,7 +80,7 @@ def testNonFinite(self):
  trials = [t1, t2]
  exptr.evaluate(trials)
  for trial in trials:
- self.assertEmpty(trial.final_measurement.metrics)
+ self.assertEmpty(trial.final_measurement_or_die.metrics)
  self.assertTrue(trial.infeasible)
 
  def testNotInSearchSpace(self):

vizier/_src/benchmarks/experimenters/shifting_experimenter_test.py

@@ -68,8 +68,9 @@ def test_numpy_experimenter(self, func):
  metric_name = exptr.problem_statement().metric_information.item().name
 
  self.assertAlmostEqual(
- t_shifted.final_measurement.metrics[metric_name].value,
- t.final_measurement.metrics[metric_name].value)
+ t_shifted.final_measurement_or_die.metrics[metric_name].value,
+ t.final_measurement_or_die.metrics[metric_name].value,
+ )
  self.assertEqual(t.status, t_shifted.status)
 
  # Check parameter bounds are shifted.
@@ -106,8 +107,9 @@ def test_evaluate_shift(self):
  metric_name = exptr.problem_statement().metric_information.item().name
 
  self.assertAlmostEqual(
- t_shifted.final_measurement.metrics[metric_name].value,
- t.final_measurement.metrics[metric_name].value)
+ t_shifted.final_measurement_or_die.metrics[metric_name].value,
+ t.final_measurement_or_die.metrics[metric_name].value,
+ )
  self.assertEqual(t.status, t_shifted.status)
  self.assertNotEqual(t.parameters, t_shifted.parameters)
 

vizier/_src/benchmarks/experimenters/sign_flip_experimenter_test.py

@@ -55,28 +55,32 @@ def test_flipped_bbob(self, flip_objectives_only: bool):
  flipped_exptr.evaluate([suggestion_for_flipped])
 
  self.assertEqual(
- suggestion_for_original.final_measurement.metrics[metric_name].value,
+ suggestion_for_original.final_measurement_or_die.metrics[
+ metric_name
+ ].value,
  -1.0
- * suggestion_for_flipped.final_measurement.metrics[metric_name].value,
+ * suggestion_for_flipped.final_measurement_or_die.metrics[
+ metric_name
+ ].value,
  )
 
  aux_metric_name = metric_name + '_before_noise'
  if flip_objectives_only:
  self.assertEqual(
- suggestion_for_original.final_measurement.metrics[
+ suggestion_for_original.final_measurement_or_die.metrics[
  aux_metric_name
  ].value,
- suggestion_for_flipped.final_measurement.metrics[
+ suggestion_for_flipped.final_measurement_or_die.metrics[
  aux_metric_name
  ].value,
  )
  else:
  self.assertEqual(
- suggestion_for_original.final_measurement.metrics[
+ suggestion_for_original.final_measurement_or_die.metrics[
  aux_metric_name
  ].value,
  -1.0
- * suggestion_for_flipped.final_measurement.metrics[
+ * suggestion_for_flipped.final_measurement_or_die.metrics[
  aux_metric_name
  ].value,
  )

vizier/_src/benchmarks/experimenters/sparse_experimenter_test.py

@@ -53,7 +53,9 @@ def test_sparse_problem_statement_and_evaluate(self, sparse_param_config):
  trial.parameters[pc.name] = 2.0
  # Test that the evaluation uses only the non-sparse parameters.
  experimenter.evaluate([trial])
- self.assertEqual(trial.final_measurement.metrics['bbob_eval'].value, 8.0)
+ self.assertEqual(
+ trial.final_measurement_or_die.metrics['bbob_eval'].value, 8.0
+ )
  # Test that the evaluated trial parameters remained the same.
  self.assertLen(trial.parameters, 10)
  for trial_param_name, exptr_param_config in zip(
@@ -106,7 +108,9 @@ def _count_param_by_type(problem):
  trial.parameters[pc.name] = 2.0
  # Test that the evaluation uses only the non-sparse parameters.
  experimenter.evaluate([trial])
- self.assertEqual(trial.final_measurement.metrics['bbob_eval'].value, 8.0)
+ self.assertEqual(
+ trial.final_measurement_or_die.metrics['bbob_eval'].value, 8.0
+ )
  # Test that the evaluated trial parameters remained the same.
  total_param_count = 2 + int_count + cat_count + discrete_count + float_count
  self.assertLen(trial.parameters, total_param_count)

vizier/_src/benchmarks/experimenters/surrogate_experimenter_test.py

@@ -65,7 +65,7 @@ def test_e2e(self):
  for trial in trials:
  self.assertEqual(trial.status, vz.TrialStatus.COMPLETED)
  self.assertContainsSubset(
- trial.final_measurement.metrics.keys(), ['metric']
+ trial.final_measurement_or_die.metrics.keys(), ['metric']
  )