From da3b051bdc9da391a6735ef75330f3a5c548abb5 Mon Sep 17 00:00:00 2001 From: You Quan Chong Date: Wed, 30 Sep 2020 14:34:14 -0700 Subject: [PATCH] Cleanup studies after integration test PiperOrigin-RevId: 334680911 --- .../tuner/optimizer_client.py | 4 +++- .../integration/tuner_integration_test.py | 20 +++++++++++++++++++ .../tuner/tests/unit/tuner_test.py | 11 +++++----- src/python/tensorflow_cloud/tuner/tuner.py | 2 +- 4 files changed, 29 insertions(+), 8 deletions(-) diff --git a/src/python/tensorflow_cloud/tuner/optimizer_client.py b/src/python/tensorflow_cloud/tuner/optimizer_client.py index e7e83e35..fdcee999 100644 --- a/src/python/tensorflow_cloud/tuner/optimizer_client.py +++ b/src/python/tensorflow_cloud/tuner/optimizer_client.py @@ -290,6 +290,7 @@ def delete_study(self, study_name: Text = None) -> None: .format(study_name)) tf.get_logger().info("DeleteStudy failed.") raise + tf.get_logger().info("Study deleted: {}.".format(study_name)) def _obtain_long_running_operation(self, resp): """Obtain the long-running operation.""" @@ -467,8 +468,9 @@ def _get_study( study_should_exist: Indicates whether it should be assumed that the study with the given study_id exists. """ - tf.get_logger().info("Study already exists. Load existing study...") study_name = "{}/studies/{}".format(study_parent, study_id) + tf.get_logger().info( + "Study already exists: {}.\nLoad existing study...".format(study_name)) num_tries = 0 while True: try: diff --git a/src/python/tensorflow_cloud/tuner/tests/integration/tuner_integration_test.py b/src/python/tensorflow_cloud/tuner/tests/integration/tuner_integration_test.py index 77bf1216..00700e67 100644 --- a/src/python/tensorflow_cloud/tuner/tests/integration/tuner_integration_test.py +++ b/src/python/tensorflow_cloud/tuner/tests/integration/tuner_integration_test.py @@ -24,6 +24,7 @@ import tensorflow as tf from tensorflow import keras from tensorflow_cloud import CloudTuner +from tensorflow_cloud.tuner import optimizer_client # If input dataset is created outside tuner.search(), # it requires eager execution even in TF 1.x. @@ -123,6 +124,10 @@ def _dist_search_fn_wrapper(args): class _CloudTunerIntegrationTestBase(tf.test.TestCase): + def setUp(self): + super(_CloudTunerIntegrationTestBase, self).setUp() + self._study_id = None + def _assert_output(self, fn, regex_str): stdout = io.StringIO() with contextlib.redirect_stdout(stdout): @@ -136,6 +141,13 @@ def _assert_results_summary(self, fn): def tearDown(self): super(_CloudTunerIntegrationTestBase, self).tearDown() + + # Delete the study used in the test, if present + if self._study_id: + service = optimizer_client.create_or_load_study( + _PROJECT_ID, _REGION, self._study_id, None) + service.delete_study() + tf.keras.backend.clear_session() @@ -149,6 +161,8 @@ def setUp(self): def testCloudTunerHyperparameters(self): """Test case to configure Tuner with HyperParameters object.""" study_id = "{}_hyperparameters".format(_STUDY_ID_BASE) + self._study_id = study_id + tuner = CloudTuner( _build_model, project_id=_PROJECT_ID, @@ -196,6 +210,8 @@ def testCloudTunerDatasets(self): ) study_id = "{}_dataset".format(_STUDY_ID_BASE) + self._study_id = study_id + tuner = CloudTuner( _build_model, project_id=_PROJECT_ID, @@ -253,6 +269,8 @@ def testCloudTunerStudyConfig(self): } study_id = "{}_study_config".format(_STUDY_ID_BASE) + self._study_id = study_id + tuner = CloudTuner( _build_model, project_id=_PROJECT_ID, @@ -286,6 +304,7 @@ class CloudTunerInDistributedIntegrationTest(_CloudTunerIntegrationTestBase): def testCloudTunerInProcessDistributedTuning(self): """Test case to simulate multiple parallel tuning workers.""" study_id = "{}_dist".format(_STUDY_ID_BASE) + self._study_id = study_id with multiprocessing.Pool(processes=_NUM_PARALLEL_TRIALS) as pool: results = pool.map( @@ -300,6 +319,7 @@ def testCloudTunerInProcessDistributedTuning(self): def testCloudTunerAIPlatformTrainingDistributedTuning(self): """Test case of parallel tuning using CAIP Training as flock manager.""" + # TODO(b/169697464): Implement test for tuning with CAIP Training study_id = "{}_caip_dist".format(_STUDY_ID_BASE) del study_id diff --git a/src/python/tensorflow_cloud/tuner/tests/unit/tuner_test.py b/src/python/tensorflow_cloud/tuner/tests/unit/tuner_test.py index 498cb9ef..d814b169 100644 --- a/src/python/tensorflow_cloud/tuner/tests/unit/tuner_test.py +++ b/src/python/tensorflow_cloud/tuner/tests/unit/tuner_test.py @@ -54,8 +54,7 @@ def setUp(self): self._region = "us-central1" self._project_id = "project-a" self._trial_parent = "projects/{}/locations/{}/studies/{}".format( - self._project_id, self._region, - "CloudTuner_study_{}".format(self._study_id) + self._project_id, self._region, self._study_id ) self._container_uri = "test_container_uri", hps = hp_module.HyperParameters() @@ -138,7 +137,7 @@ def test_tuner_initialization_with_hparams(self): (self.mock_optimizer_client_module.create_or_load_study .assert_called_with(self._project_id, self._region, - "CloudTuner_study_{}".format(self._study_id), + self._study_id, self._study_config)) def test_tuner_initialization_with_study_config(self): @@ -146,7 +145,7 @@ def test_tuner_initialization_with_study_config(self): (self.mock_optimizer_client_module.create_or_load_study .assert_called_with(self._project_id, self._region, - "CloudTuner_study_{}".format(self._study_id), + self._study_id, self._study_config)) def test_remote_tuner_initialization_with_study_config(self): @@ -154,7 +153,7 @@ def test_remote_tuner_initialization_with_study_config(self): (self.mock_optimizer_client_module.create_or_load_study .assert_called_with(self._project_id, self._region, - "CloudTuner_study_{}".format(self._study_id), + self._study_id, self._study_config)) def test_tuner_initialization_neither_hparam_nor_study_config(self): @@ -178,7 +177,7 @@ def test_tuner_initialization_with_study_config_and_max_trials(self): (self.mock_optimizer_client_module.create_or_load_study .assert_called_with(self._project_id, self._region, - "CloudTuner_study_{}".format(self._study_id), + self._study_id, self._study_config)) def test_create_trial_initially(self): diff --git a/src/python/tensorflow_cloud/tuner/tuner.py b/src/python/tensorflow_cloud/tuner/tuner.py index 65533ef1..a630009b 100644 --- a/src/python/tensorflow_cloud/tuner/tuner.py +++ b/src/python/tensorflow_cloud/tuner/tuner.py @@ -113,7 +113,7 @@ def __init__( self.max_trials = max_trials if study_id: - self.study_id = "CloudTuner_study_{}".format(study_id) + self.study_id = study_id else: self.study_id = "CloudTuner_study_{}".format( datetime.datetime.now().strftime("%Y%m%d_%H%M%S")