Skip to content

Introduce UPSERT in set_trial_system_attr#5741

Merged
HideakiImamura merged 2 commits intooptuna:masterfrom
porink0424:fix/introduce-upsert-in-set_trial_system_attr
Nov 8, 2024
Merged

Introduce UPSERT in set_trial_system_attr#5741
HideakiImamura merged 2 commits intooptuna:masterfrom
porink0424:fix/introduce-upsert-in-set_trial_system_attr

Conversation

@porink0424
Copy link
Copy Markdown
Member

@porink0424 porink0424 commented Nov 1, 2024 •
edited
Loading

Motivation

This PR is a follow-up to #5703. This introduces UPSERT to set_trial_system_attr to improve the generated SQL queries.

Description of the changes

  • The logic of _set_trial_user_attr_without_commit and _set_trial_system_attr_without_commit is almost identical, so they have been consolidated into a single function, _set_trial_attr_without_commit, with introducing a model_cls argument. Note that, as with Introduce UPSERT in set_trial_user_attr #5703, PostgreSQL is still not supported.

Verification of SQL changes

MySQL

import optuna
import logging

optuna.logging.set_verbosity(optuna.logging.WARNING)

study = optuna.create_study(storage="mysql+pymysql://optuna:password@127.0.0.1:3306/optuna")
storage = study._storage


def objective(trial: optuna.Trial) -> float:
    storage.set_trial_system_attr(trial._trial_id, "system_attr_1", "value_1")

    logging.basicConfig()
    logging.getLogger("sqlalchemy.engine").setLevel(logging.INFO)
    storage.set_trial_system_attr(trial._trial_id, "system_attr_1", "new_value_1")
    storage.set_trial_system_attr(trial._trial_id, "system_attr_2", "value_2")
    logging.getLogger("sqlalchemy.engine").setLevel(logging.WARNING)

    return trial.suggest_float("x", 0.0, 1.0) ** 2


study.optimize(objective, n_trials=1)

Before

...
INFO:sqlalchemy.engine.Engine:SELECT trial_system_attributes.trial_system_attribute_id AS trial_system_attributes_trial_system_attribute_id, trial_system_attributes.trial_id AS trial_system_attributes_trial_id, trial_system_attributes.`key` AS trial_system_attributes_key, trial_system_attributes.value_json AS trial_system_attributes_value_json 
FROM trial_system_attributes 
WHERE trial_system_attributes.trial_id = %(trial_id_1)s AND trial_system_attributes.`key` = %(key_1)s
INFO:sqlalchemy.engine.Engine:[cached since 0.00552s ago] {'trial_id_1': 1520, 'key_1': 'system_attr_1'}
INFO:sqlalchemy.engine.Engine:UPDATE trial_system_attributes SET value_json=%(value_json)s WHERE trial_system_attributes.trial_system_attribute_id = %(trial_system_attributes_trial_system_attribute_id)s
INFO:sqlalchemy.engine.Engine:[generated in 0.00014s] {'value_json': '"new_value_1"', 'trial_system_attributes_trial_system_attribute_id': 13167}
...
INFO:sqlalchemy.engine.Engine:SELECT trial_system_attributes.trial_system_attribute_id AS trial_system_attributes_trial_system_attribute_id, trial_system_attributes.trial_id AS trial_system_attributes_trial_id, trial_system_attributes.`key` AS trial_system_attributes_key, trial_system_attributes.value_json AS trial_system_attributes_value_json 
FROM trial_system_attributes 
WHERE trial_system_attributes.trial_id = %(trial_id_1)s AND trial_system_attributes.`key` = %(key_1)s
INFO:sqlalchemy.engine.Engine:[cached since 0.01052s ago] {'trial_id_1': 1520, 'key_1': 'system_attr_2'}
INFO:sqlalchemy.engine.Engine:INSERT INTO trial_system_attributes (trial_id, `key`, value_json) VALUES (%(trial_id)s, %(key)s, %(value_json)s)
INFO:sqlalchemy.engine.Engine:[cached since 0.01002s ago] {'trial_id': 1520, 'key': 'system_attr_2', 'value_json': '"value_2"'}
...

After

...
INFO:sqlalchemy.engine.Engine:INSERT INTO trial_system_attributes (trial_id, `key`, value_json) VALUES (%(trial_id)s, %(key)s, %(value_json)s) ON DUPLICATE KEY UPDATE value_json = VALUES(value_json)
INFO:sqlalchemy.engine.Engine:[no key 0.00009s] {'trial_id': 1521, 'key': 'system_attr_1', 'value_json': '"new_value_1"'}
...
INFO:sqlalchemy.engine.Engine:INSERT INTO trial_system_attributes (trial_id, `key`, value_json) VALUES (%(trial_id)s, %(key)s, %(value_json)s) ON DUPLICATE KEY UPDATE value_json = VALUES(value_json)
INFO:sqlalchemy.engine.Engine:[no key 0.00008s] {'trial_id': 1521, 'key': 'system_attr_2', 'value_json': '"value_2"'}
...

sqlite

...
study = optuna.create_study(storage="sqlite:///example.db")
...

Before

...
INFO:sqlalchemy.engine.Engine:SELECT trial_system_attributes.trial_system_attribute_id AS trial_system_attributes_trial_system_attribute_id, trial_system_attributes.trial_id AS trial_system_attributes_trial_id, trial_system_attributes."key" AS trial_system_attributes_key, trial_system_attributes.value_json AS trial_system_attributes_value_json 
FROM trial_system_attributes 
WHERE trial_system_attributes.trial_id = ? AND trial_system_attributes."key" = ?
INFO:sqlalchemy.engine.Engine:[cached since 0.003618s ago] (1, 'system_attr_1')
INFO:sqlalchemy.engine.Engine:UPDATE trial_system_attributes SET value_json=? WHERE trial_system_attributes.trial_system_attribute_id = ?
INFO:sqlalchemy.engine.Engine:[generated in 0.00009s] ('"new_value_1"', 1)
...
INFO:sqlalchemy.engine.Engine:SELECT trial_system_attributes.trial_system_attribute_id AS trial_system_attributes_trial_system_attribute_id, trial_system_attributes.trial_id AS trial_system_attributes_trial_id, trial_system_attributes."key" AS trial_system_attributes_key, trial_system_attributes.value_json AS trial_system_attributes_value_json 
FROM trial_system_attributes 
WHERE trial_system_attributes.trial_id = ? AND trial_system_attributes."key" = ?
INFO:sqlalchemy.engine.Engine:[cached since 0.006849s ago] (1, 'system_attr_2')
INFO:sqlalchemy.engine.Engine:INSERT INTO trial_system_attributes (trial_id, "key", value_json) VALUES (?, ?, ?)
INFO:sqlalchemy.engine.Engine:[cached since 0.006491s ago] (1, 'system_attr_2', '"value_2"')
...

After

...
INFO:sqlalchemy.engine.Engine:INSERT INTO trial_system_attributes (trial_id, "key", value_json) VALUES (?, ?, ?) ON CONFLICT (trial_id, "key") DO UPDATE SET value_json = excluded.value_json
INFO:sqlalchemy.engine.Engine:[no key 0.00009s] (2, 'system_attr_1', '"new_value_1"')
...
INFO:sqlalchemy.engine.Engine:INSERT INTO trial_system_attributes (trial_id, "key", value_json) VALUES (?, ?, ?) ON CONFLICT (trial_id, "key") DO UPDATE SET value_json = excluded.value_json
INFO:sqlalchemy.engine.Engine:[no key 0.00009s] (2, 'system_attr_2', '"value_2"')
...

@porink0424 porink0424 marked this pull request as ready for review November 1, 2024 07:14
@c-bata c-bata added the enhancement Change that does not break compatibility and not affect public interfaces, but improves performance. label Nov 1, 2024
@c-bata
Copy link
Copy Markdown
Member

c-bata commented Nov 1, 2024

@gen740 @nabenabe0928 Could you review this PR if you have time?

nabenabe0928
nabenabe0928 approved these changes Nov 7, 2024
View reviewed changes
Copy link
Copy Markdown
Contributor

@nabenabe0928 nabenabe0928 left a comment

Choose a reason for hiding this comment

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

LGTM, but I left a nit-picky comment

HideakiImamura
HideakiImamura approved these changes Nov 8, 2024
View reviewed changes
Copy link
Copy Markdown
Member

@HideakiImamura HideakiImamura left a comment

Choose a reason for hiding this comment

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

Looks very good.

@HideakiImamura HideakiImamura merged commit d9327ca into optuna:master Nov 8, 2024
@HideakiImamura HideakiImamura added this to the v4.1.0 milestone Nov 8, 2024
@porink0424 porink0424 deleted the fix/introduce-upsert-in-set_trial_system_attr branch November 20, 2024 01:25
@gen740 gen740 removed their assignment Jun 3, 2025
@jaikumarm jaikumarm mentioned this pull request Sep 18, 2025
Merged
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

@HideakiImamura HideakiImamura HideakiImamura approved these changes

+1 more reviewer

@nabenabe0928 nabenabe0928 nabenabe0928 approved these changes

Reviewers whose approvals may not affect merge requirements

Assignees

@nabenabe0928 nabenabe0928

Labels

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

Projects

None yet

Milestone

v4.1.0

Development

Successfully merging this pull request may close these issues.

5 participants

@porink0424 @c-bata @HideakiImamura @nabenabe0928 @gen740