from optuna import create_study
from optuna.storages import GrpcStorageProxy
from concurrent.futures import ThreadPoolExecutor
from concurrent.futures import as_completed

num_enqueued = 50
study = create_study(
    storage=GrpcStorageProxy(
        host="localhost",
        port=13000,
    )
)
for i in range(num_enqueued):
    study.enqueue_trial({"i": i})

trial_id_set = set()

with ThreadPoolExecutor(100) as pool:
    futures = []
    for i in range(num_enqueued):
        future = pool.submit(study._pop_waiting_trial_id)
        futures.append(future)

    for future in as_completed(futures):
        trial_id_set.add(future.result())
print(f"{num_enqueued = },  {len(trial_id_set) = }")
assert len(trial_id_set) == num_enqueued