Skip to content

Cache distributions to skip consistency check#6301

Merged
y0z merged 4 commits intooptuna:masterfrom
not522:cache-dist
Nov 5, 2025
Merged

Cache distributions to skip consistency check#6301
y0z merged 4 commits intooptuna:masterfrom
not522:cache-dist

Conversation

@not522
Copy link
Copy Markdown
Member

@not522 not522 commented Oct 14, 2025 •
edited
Loading

Motivation

RDBStorage performs consistency checks with past distributions in set_trial_param, but it can be skipped by caching past distributions in _CachedStorage.
In the original _CachedStorage (#1140), it was implemented, but lost that feature in #4631 and removed in #5978. This PR reintroduces that functionality.

Description of the changes

  • Cache distributions in _CachedStorage to skip consistency check

Benchmark

import optuna

def objective(trial: optuna.Trial) -> float:
    x = trial.suggest_float("x", -100, 100)
    y = trial.suggest_int("y", -100, 100)
    return x**2 + y**2

sampler = optuna.samplers.TPESampler(seed=42, multivariate=True, constant_liar=True)
study = optuna.create_study(sampler=sampler, storage="sqlite:///tmp.db")
study.optimize(objective, n_trials=1000)
master PR
10.164s 9.806s

@not522 not522 added the enhancement Change that does not break compatibility and not affect public interfaces, but improves performance. label Oct 14, 2025
@not522 not522 marked this pull request as draft October 14, 2025 08:39
not522
not522 commented Oct 17, 2025
View reviewed changes
@not522 not522 marked this pull request as ready for review October 17, 2025 08:15
@not522
Copy link
Copy Markdown
Member Author

not522 commented Oct 22, 2025 •
edited
Loading

The improvement from this PR is proportional to the number of times set_trial_param is called, so when the number of parameters is large, the difference becomes even more pronounced as follows:

import optuna

def objective(trial: optuna.Trial) -> float:
    return sum(trial.suggest_float(f"x{i}", -100, 100) ** 2 for i in range(10))

sampler = optuna.samplers.RandomSampler(seed=42)
study = optuna.create_study(sampler=sampler, storage="sqlite:///tmp.db")
study.optimize(objective, n_trials=1000)
master PR
11.885s 10.554s

@y0z
Copy link
Copy Markdown
Member

y0z commented Oct 30, 2025

@sawa3030 @kAIto47802 Could you review this PR?

sawa3030
sawa3030 approved these changes Oct 30, 2025
View reviewed changes
Copy link
Copy Markdown
Collaborator

@sawa3030 sawa3030 left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thank you for the additional benchmarking. I have reviewed and confirmed the changes. Since the caching algorithm is relatively simple here, I am in favor of merging this PR.

kAIto47802
kAIto47802 reviewed Oct 31, 2025
View reviewed changes
Copy link
Copy Markdown
Collaborator

@kAIto47802 kAIto47802 left a comment •
edited
Loading

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thank you for the PR!
It basically LGTM, with only a minor comment :octocat:

optuna/storages/_rdb/storage.py Outdated
@@ -587,25 +587,60 @@ def set_trial_param(
session, trial_id, param_name, param_value_internal, distribution
)
Copy link
Copy Markdown
Collaborator

@kAIto47802 kAIto47802 Oct 31, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

How about using _set_trial_params here, since set_trial_params performs the same operation while just hiding the previous_distribution argument from the user?

        return self._set_trial_param(trial_id, param_name, param_value_internal, distribution)

Copy link
Copy Markdown
Member Author

@not522 not522 Nov 5, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thank you. I've updated as suggested. PTAL.

@c-bata c-bata added this to the v4.7.0 milestone Nov 4, 2025
kAIto47802
kAIto47802 approved these changes Nov 5, 2025
View reviewed changes
Copy link
Copy Markdown
Collaborator

@kAIto47802 kAIto47802 left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thank you for the update! LGTM

@y0z y0z merged commit 9cdcaf3 into optuna:master Nov 5, 2025
12 checks passed
@not522 not522 deleted the cache-dist branch November 5, 2025 21:38
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

@kAIto47802 kAIto47802 kAIto47802 approved these changes

@sawa3030 sawa3030 sawa3030 approved these changes

Assignees

@kAIto47802 kAIto47802

@sawa3030 sawa3030

Labels

enhancement Change that does not break compatibility and not affect public interfaces, but improves performance.

Projects

None yet

Milestone

v4.7.0

Development

Successfully merging this pull request may close these issues.

5 participants

@not522 @y0z @kAIto47802 @sawa3030 @c-bata