fix tests

keras-team · Oct 10, 2023 · 40cd10d · 40cd10d
1 parent 2b07de2
commit 40cd10d
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 3 deletions.
diff --git a/keras_tuner/distribute/oracle_client_test.py b/keras_tuner/distribute/oracle_client_test.py
@@ -110,6 +110,7 @@ def build_model(hp):
             max_trials=10,
             directory=tmp_path,
         )
+        tuner.search(x, y, validation_data=(x, y), epochs=1, batch_size=2)
 
         # Only worker makes it to this point, server runs until thread stops.
         assert dist_utils.has_chief_oracle()
@@ -118,8 +119,6 @@ def build_model(hp):
             tuner.oracle, keras_tuner.distribute.oracle_client.OracleClient
         )
 
-        tuner.search(x, y, validation_data=(x, y), epochs=1, batch_size=2)
-
         # Suppress warnings about optimizer state not being restored by
         # tf.keras.
 

diff --git a/keras_tuner/engine/base_tuner.py b/keras_tuner/engine/base_tuner.py
@@ -132,7 +132,7 @@ def __init__(
             self._populate_initial_space()
 
         # Run in distributed mode.
-        if dist_utils.has_chief_oracle():
+        if dist_utils.has_chief_oracle() and not dist_utils.is_chief_oracle():
             # Proxies requests to the chief oracle.
             # Avoid import at the top, to avoid inconsistent protobuf versions.
             from keras_tuner.distribute import oracle_client
@@ -216,6 +216,7 @@ def search(self, *fit_args, **fit_kwargs):
             from keras_tuner.distribute import oracle_chief
 
             oracle_chief.start_server(self.oracle)
+            return
 
         self.on_search_begin()
         while True: