Merge pull request optuna#6004 from fusawa-yugo/fusawa-yugo/survey-error-in-grpc-server

Fix a bug that a gRPC server doesn't work with JournalStorage

Author: nabenabe0928

optuna/storages/_grpc/servicer.py

@@ -249,7 +249,7 @@ def SetTrialStateValues(
     ) -> api_pb2.SetTrialStateValuesReply:
         trial_id = request.trial_id
         state = request.state
-        values = request.values
+        values = list(request.values)
         try:
             trial_updated = self._backend.set_trial_state_values(
                 trial_id, _from_proto_trial_state(state), values

optuna/testing/storages.py

@@ -30,7 +30,8 @@
     "cached_sqlite",
     "journal",
     "journal_redis",
-    "grpc",
+    "grpc_rdb",
+    "grpc_journal_file",
 ]
 
 
@@ -84,12 +85,34 @@ def __enter__(
                 "redis", fakeredis.FakeStrictRedis()  # type: ignore[no-untyped-call]
             )
             return optuna.storages.JournalStorage(journal_redis_storage)
+        elif self.storage_specifier == "grpc_journal_file":
+            self.tempfile = self.extra_args.get("file", NamedTemporaryFilePool().tempfile())
+            assert self.tempfile is not None
+            port = _find_free_port()
+            storage = optuna.storages.JournalStorage(
+                optuna.storages.journal.JournalFileBackend(self.tempfile.name)
+            )
+
+            self.server = optuna.storages._grpc.server.make_server(storage, "localhost", port)
+            self.thread = threading.Thread(target=self.server.start)
+            self.thread.start()
+
+            proxy = GrpcStorageProxy(host="localhost", port=port)
+
+            # Wait until the server is ready.
+            while True:
+                try:
+                    proxy.get_all_studies()
+                    return proxy
+                except grpc.RpcError:
+                    time.sleep(1)
+                    continue
         elif "journal" in self.storage_specifier:
             self.tempfile = self.extra_args.get("file", NamedTemporaryFilePool().tempfile())
             assert self.tempfile is not None
             file_storage = JournalFileBackend(self.tempfile.name)
             return optuna.storages.JournalStorage(file_storage)
-        elif self.storage_specifier == "grpc":
+        elif self.storage_specifier == "grpc_rdb":
             self.tempfile = NamedTemporaryFilePool().tempfile()
             url = "sqlite:///{}".format(self.tempfile.name)
             port = _find_free_port()

tests/storages_tests/test_storages.py

@@ -773,7 +773,7 @@ def test_get_all_trials(storage_mode: str) -> None:
 def test_get_all_trials_params_order(storage_mode: str, param_names: list[str]) -> None:
     # We don't actually require that all storages to preserve the order of parameters,
     # but all current implementations except for GrpcStorageProxy do, so we test this property.
-    if storage_mode == "grpc":
+    if storage_mode in ("grpc_rdb", "grpc_journal_file"):
         pytest.skip("GrpcStorageProxy does not preserve the order of parameters.")
 
     with StorageSupplier(storage_mode) as storage:

tests/study_tests/test_study.py

@@ -1036,7 +1036,7 @@ def terminate_study(study: Study, trial: FrozenTrial) -> None:
 
 @pytest.mark.parametrize("storage_mode", STORAGE_MODES)
 def test_get_trials(storage_mode: str) -> None:
-    if storage_mode == "grpc":
+    if storage_mode in ("grpc_rdb", "grpc_journal_file"):
         pytest.skip("gRPC storage doesn't use `copy.deepcopy`.")
 
     with StorageSupplier(storage_mode) as storage:
@@ -1651,7 +1651,7 @@ def test_tell_from_another_process() -> None:
 
 @pytest.mark.parametrize("storage_mode", STORAGE_MODES)
 def test_pop_waiting_trial_thread_safe(storage_mode: str) -> None:
-    if "sqlite" == storage_mode or "cached_sqlite" == storage_mode or "grpc" == storage_mode:
+    if storage_mode in ("sqlite", "cached_sqlite", "grpc_rdb"):
         pytest.skip("study._pop_waiting_trial is not thread-safe on SQLite3")
 
     num_enqueued = 10