Remove overridden get_trial, make sure metrics of trials are updated locally, and

make sure trial.best_step is an integer

PiperOrigin-RevId: 343406054

Author: chongyouquan

src/python/tensorflow_cloud/tuner/tests/unit/optimizer_client_test.py

@@ -393,6 +393,17 @@ def test_get_trial(self):
         mock_get_trial.assert_called_once_with(name=self._trial_name)
         self.assertEqual(trial, expected_trial)
 
+    def test_get_trial_with_404_raises(self):
+        mock_get_trial = mock.MagicMock()
+        mock_get_trial.return_value.execute.side_effect = errors.HttpError(
+            httplib2.Response(info={"status": 404}), b"")
+
+        self._mock_discovery.projects().locations().studies().trials(
+            ).get = mock_get_trial
+
+        with self.assertRaises(errors.HttpError):
+            self._client.get_trial(trial_id="1")
+
     def test_list_trials(self):
         mock_list_trials = mock.MagicMock()
         expected_trials = {

src/python/tensorflow_cloud/tuner/tests/unit/tuner_test.py

@@ -87,7 +87,7 @@ def setUp(self):
         self._test_trial = trial_module.Trial(
             hyperparameters=self._test_hyperparameters,
             trial_id="1",
-            status=trial_module.TrialStatus,
+            status=trial_module.TrialStatus.RUNNING,
         )
         # TODO(b/170687807) Switch from using "{}".format() to f-string
         self._job_id = "{}_{}".format(self._study_id, self._test_trial.trial_id)
@@ -253,7 +253,8 @@ def test_create_trial_after_early_stopping(self):
         self.assertEqual(trial.hyperparameters.values, {})
         self.assertEqual(trial.status, trial_module.TrialStatus.STOPPED)
 
-    def test_update_trial(self):
+    @mock.patch.object(oracle_module.Oracle, "update_trial", auto_spec=True)
+    def test_update_trial(self, mock_super_update_trial):
         self._tuner_with_hparams()
 
         self.mock_client.should_trial_stop.return_value = True
@@ -277,6 +278,9 @@ def test_update_trial(self):
         )
         self.mock_client.should_trial_stop.assert_called_once_with("1")
         self.assertEqual(status, trial_module.TrialStatus.STOPPED)
+        mock_super_update_trial.assert_called_once_with(
+            "1", {"val_acc": 0.8}, 3
+        )
 
     def test_end_trial_success(self):
         self._tuner_with_hparams()
@@ -285,17 +289,29 @@ def test_end_trial_success(self):
             "state": "COMPLETED",
             "parameters": [{"parameter": "learning_rate", "floatValue": 0.01}],
             "finalMeasurement": {
-                "stepCount": 3,
+                "stepCount": "3",
                 "metrics": [{"metric": "val_acc", "value": 0.7}],
             },
             "trial_infeasible": False,
             "infeasible_reason": None,
         }
-
+        mock_save_trial = mock.Mock()
+        self.tuner.oracle._save_trial = mock_save_trial
         self.tuner.oracle.ongoing_trials = {"tuner_0": self._test_trial}
+        expected_trial = trial_module.Trial(
+            hyperparameters=self._test_hyperparameters,
+            trial_id="1",
+            status=trial_module.TrialStatus.COMPLETED,
+        )
+        expected_trial.best_step = 3
+        expected_trial.score = 0.7
+
         self.tuner.oracle.end_trial(trial_id="1")
+
         self.mock_client.complete_trial.assert_called_once_with(
             "1", False, None)
+        self.assertEqual(repr(mock_save_trial.call_args[0][0].get_state()),
+                         repr(expected_trial.get_state()))
 
     def test_end_trial_infeasible_trial(self):
         self._tuner_with_hparams()
@@ -319,35 +335,6 @@ def test_end_trial_invalid_status(self):
         with self.assertRaises(ValueError):
             self.tuner.oracle.end_trial(trial_id="1", status="FOO")
 
-    def test_get_trial_success(self):
-        self._tuner_with_hparams()
-        self.mock_client.get_trial.return_value = {
-            "name": "1",
-            "state": "COMPLETED",
-            "parameters": [{"parameter": "learning_rate", "floatValue": 0.01}],
-            "finalMeasurement": {
-                "stepCount": 3,
-                "metrics": [{"metric": "val_acc", "value": 0.7}],
-            },
-            "trial_infeasible": False,
-            "infeasible_reason": None,
-        }
-        trial = self.tuner.oracle.get_trial(trial_id="1")
-        self.mock_client.get_trial.assert_called_once_with("1")
-        self.assertEqual(trial.trial_id, "1")
-        self.assertEqual(trial.score, 0.7)
-        self.assertEqual(trial.status, trial_module.TrialStatus.COMPLETED)
-        self.assertEqual(trial.hyperparameters.values, {"learning_rate": 0.01})
-
-    def test_get_trial_failed(self):
-        self._tuner_with_hparams()
-        self.mock_client.get_trial.return_value = {
-            "name": "1",
-            "state": "FOO"
-        }
-        with self.assertRaises(ValueError):
-            self.tuner.oracle.get_trial(trial_id="1")
-
     def test_get_best_trials(self):
         self._tuner_with_hparams()
 
@@ -358,7 +345,7 @@ def test_get_best_trials(self):
                 "parameters":
                     [{"parameter": "learning_rate", "floatValue": 0.01}],
                 "finalMeasurement": {
-                    "stepCount": 3,
+                    "stepCount": "3",
                     "metrics": [{"metric": "val_acc", "value": 0.7}],
                 },
                 "trial_infeasible": False,
@@ -370,7 +357,7 @@ def test_get_best_trials(self):
                 "parameters":
                     [{"parameter": "learning_rate", "floatValue": 0.001}],
                 "finalMeasurement": {
-                    "stepCount": 3,
+                    "stepCount": "3",
                     "metrics": [{"metric": "val_acc", "value": 0.9}],
                 },
                 "trial_infeasible": False,
@@ -425,7 +412,7 @@ def test_get_best_trials_multi_tuners(self):
                 "parameters":
                     [{"parameter": "learning_rate", "floatValue": 0.01}],
                 "finalMeasurement": {
-                    "stepCount": 3,
+                    "stepCount": "3",
                     "metrics": [{"metric": "val_acc", "value": 0.7}],
                 },
                 "trial_infeasible": False,
@@ -437,7 +424,7 @@ def test_get_best_trials_multi_tuners(self):
                 "parameters":
                     [{"parameter": "learning_rate", "floatValue": 0.001}],
                 "finalMeasurement": {
-                    "stepCount": 3,
+                    "stepCount": "3",
                     "metrics": [{"metric": "val_acc", "value": 0.9}],
                 },
                 "trial_infeasible": False,
@@ -458,6 +445,11 @@ def test_get_best_trials_multi_tuners(self):
         self.assertEqual(best_trials_1[0].score, 0.9)
         self.assertEqual(best_trials_1[0].best_step, 3)
 
+    def test_get_single_objective(self):
+        self._tuner_with_hparams()
+        self.assertEqual([self.tuner.oracle.objective],
+                         self.tuner.oracle._get_objective())
+
     @mock.patch.object(super_tuner.Tuner, "__init__", auto_spec=True)
     @mock.patch.object(tf.summary, "create_file_writer", auto_spec=True)
     @mock.patch.object(hparams_api, "hparams", auto_spec=True)

src/python/tensorflow_cloud/tuner/tests/unit/utils_test.py

@@ -189,7 +189,7 @@
         {"parameter": "learning_rate", "floatValue": 0.0001},
     ],
     "finalMeasurement": {
-        "stepCount": 1,
+        "stepCount": "1",
         "metrics": [{"value": 0.9}],
     },
 }
@@ -305,13 +305,14 @@ def test_convert_optimizer_trial_to_hps(self):
         trial_hps = utils.convert_optimizer_trial_to_hps(hps, OPTIMIZER_TRIAL)
         self.assertDictEqual(trial_hps.values, EXPECTED_TRIAL_HPS)
 
-    def test_convert_optimizer_trial_to_keras_trial(self):
+    def test_convert_completed_optimizer_trial_to_keras_trial(self):
         hps = hp_module.HyperParameters()
         hps.Choice("learning_rate", [1e-4, 1e-3, 1e-2])
-        trial = utils.convert_optimizer_trial_to_keras_trial(
+        trial = utils.convert_completed_optimizer_trial_to_keras_trial(
             COMPLETED_OPTIMIZER_TRIAL, hps)
         self.assertEqual(trial.trial_id, "trial_1")
         self.assertEqual(trial.score, 0.9)
+        self.assertEqual(trial.best_step, 1)
         self.assertEqual(trial.status, trial_module.TrialStatus.COMPLETED)
         self.assertEqual(
             trial.hyperparameters.values, {"learning_rate": 0.0001})

src/python/tensorflow_cloud/tuner/tuner.py

@@ -123,7 +123,11 @@ def __init__(
             raise ValueError('"region" is not found.')
         self._region = region
 
-        self.objective = utils.format_objective(objective)
+        # If it's just single objective, let it be an Objective instead of a
+        # list, to keep it consistent with how KerasTuner formats objectives
+        obj = utils.format_objective(objective)
+        self.objective = obj[0] if len(obj) == 1 else obj
+
         self.hyperparameters = hyperparameters
         self.max_trials = max_trials
 
@@ -218,6 +222,7 @@ def update_trial(self,
         """Used by a worker to report the status of a trial."""
         # Constructs the measurement.
         # Adds the measurement of the objective functions to a trial.
+        super(CloudOracle, self).update_trial(trial_id, metrics, step)
         elapsed_secs = time.time() - self._start_time
         if elapsed_secs < 0 or step < 0:
             raise ValueError(
@@ -227,7 +232,7 @@ def update_trial(self,
                 "At least one of {elapsed_secs, step} must be positive")
 
         metric_list = []
-        for ob in self.objective:
+        for ob in self._get_objective():
             if ob.name not in metrics:
                 tf.get_logger().info(
                     'Objective "{}" is not found in metrics.'.format(ob.name)
@@ -293,22 +298,6 @@ def end_trial(self, trial_id: Text, status: Text = "COMPLETED"):
         self._save_trial(kerastuner_trial)
         self.save()
 
-    def get_trial(self, trial_id: Text) -> trial_module.Trial:
-        """Returns a completed KerasTuner Trial given the trial_id."""
-        # Note that this is called in Tuner.on_trial_end.
-
-        optimizer_trial = self.service.get_trial(trial_id)
-
-        if optimizer_trial["state"] != "COMPLETED":
-            raise ValueError("The trial status is not COMPLETED, found {}"
-                             .format(optimizer_trial["state"]))
-
-        # Convert a completed Optimizer trial to KerasTuner Trial instance.
-        kerastuner_trial = utils.convert_optimizer_trial_to_keras_trial(
-            optimizer_trial,
-            self.hyperparameters.copy())
-        return kerastuner_trial
-
     def get_best_trials(self, num_trials: int = 1) -> List[trial_module.Trial]:
         """Returns the trials with the best objective values found so far.
 
@@ -317,14 +306,14 @@ def get_best_trials(self, num_trials: int = 1) -> List[trial_module.Trial]:
         Returns:
             List of KerasTuner Trials.
         """
-        if len(self.objective) > 1:
+        objective = self._get_objective()
+        if len(objective) > 1:
             raise ValueError(
                 "Getting the best trials for multi-objective optimization "
                 "is not supported."
             )
 
-        maximizing = (
-            utils.format_goal(self.objective[0].direction) == "MAXIMIZE")
+        maximizing = (utils.format_goal(objective[0].direction) == "MAXIMIZE")
 
         # List all trials associated with the same study
         trial_list = self.service.list_trials()
@@ -336,20 +325,26 @@ def get_best_trials(self, num_trials: int = 1) -> List[trial_module.Trial]:
 
         sorted_trials = sorted(
             optimizer_trials,
-            key=lambda t: t["finalMeasurement"]["metrics"][0]["value"],
+            key=lambda t: t["finalMeasurement"]["metrics"][0].get("value"),
             reverse=maximizing,
         )
         best_optimizer_trials = sorted_trials[:num_trials]
 
         best_trials = []
         # Convert completed Optimizer trials to KerasTuner Trial instances.
         for optimizer_trial in best_optimizer_trials:
-            kerastuner_trial = utils.convert_optimizer_trial_to_keras_trial(
-                optimizer_trial,
-                self.hyperparameters.copy())
+            kerastuner_trial = (
+                utils.convert_completed_optimizer_trial_to_keras_trial(
+                    optimizer_trial,
+                    self.hyperparameters.copy()))
             best_trials.append(kerastuner_trial)
         return best_trials
 
+    def _get_objective(self):
+        """Returns the Objective(s) as a list."""
+        return self.objective if isinstance(self.objective,
+                                            list) else [self.objective]
+
 
 class CloudTuner(tuner_module.Tuner):
     """KerasTuner interface implementation backed by CAIP Optimizer Service."""

src/python/tensorflow_cloud/tuner/utils.py

@@ -509,7 +509,7 @@ def convert_optimizer_trial_to_hps(
     return hps
 
 
-def convert_optimizer_trial_to_keras_trial(
+def convert_completed_optimizer_trial_to_keras_trial(
     optimizer_trial: Dict[Text, Any],
     hyperparameter_space: hp_module.HyperParameters,
 ) -> trial_module.Trial:
@@ -537,7 +537,7 @@ def convert_optimizer_trial_to_keras_trial(
         raise ValueError('"finalMeasurement" not found in this trial {}'
                          .format(optimizer_trial))
 
-    kerastuner_trial.best_step = final_measurement.get("stepCount", 0)
+    kerastuner_trial.best_step = int(final_measurement.get("stepCount", 0))
     kerastuner_trial.score = final_measurement["metrics"][0].get("value")
     return kerastuner_trial