Migration tests for optuna/optuna#2030

README.md

Overview

This document explains how to test migrated schema of Optuna's RDBStorage. The test consists of following two parts:

1. Test for values
2. Test for schema

Each section describes the test instruction for each database.

Setup

The following software is required:

• docker
• Python 3.6+

Then, please in stall DB drivers:

$ pip install PyMySQL psycopg2-binary

Usage

SQLite

Test for values

Install the current version of optuna and create tables:

$ pip uninstall optuna  # if necessary
$ pip install optuna
$ python migration-prepare.py sqlite:///example.db

Upgrade optuna to the patched version

$ pip uninstall optuna
$ pip install git+https://github.com/HideakiImamura/optuna.git@refactoring/change-rdb-schema-for-mo-refactoring

Check if RuntimeError occurs due to the incompatibility of the schema.

$ python migration-assert.py sqlite:///example.db
optuna version: 2.3.0
...
RuntimeError: The runtime optuna version 2.3.0 is no longer compatible with the table schema (set up by optuna 2.3.0). Please execute `$ optuna storage upgrade --storage $STORAGE_URL` for upgrading the storage.

Upgrade the schema

$ optuna storage upgrade --storage sqlite:///example.db

Check the migration results

$ python migration-assert.py sqlite:///example.db
OK

Test for schema

Please run following two scripts. The first one execute migration-prepare.py with optuna@master while the second one execute it with optuna v2.3.0 and upgrade the DB.

rm -f example.db
pip install -U git+https://github.com/optuna/optuna.git
python migration-prepare.py sqlite:///example.db
echo ".dump" | sqlite3 example.db > master.sqlite

rm -f example.db
pip install -U optuna
python migration-prepare.py sqlite:///example.db
pip install -U git+https://github.com/optuna/optuna.git
optuna storage upgrade --storage sqlite:///example.db
echo ".dump" | sqlite3 example.db > upgrade.sqlite

The schema of these two should be identical except for trials' result. Please check the diff of the dump files, for example, as follows:

python sqlite_extract_create.py master.sqlite > master-create.sqlite
python sqlite_extract_create.py upgrade.sqlite > upgrade-create.sqlite
diff master-create.sqlite upgrade-create.sqlite | less -S

MySQL

Test for values

Install the current version of optuna, launch a DB server, and create tables:

$ pip uninstall optuna  # if necessary
$ pip install optuna
$ docker run --name some-mysql -e MYSQL_ROOT_PASSWORD=test -p 3306:3306 -p 33060:33060 -d mysql:5.7
$ docker run --network host -it --rm mysql:5.7  mysql -h 127.0.0.1 -uroot -p -e "create database optunatest;"
$ python migration-prepare.py mysql+pymysql://root:test@localhost/optunatest

Upgrade optuna to the patched version

$ pip uninstall optuna
$ pip install git+https://github.com/HideakiImamura/optuna.git@refactoring/change-rdb-schema-for-mo-refactoring

Check if RuntimeError occurs due to the incompatibility of the schema.

$ python migration-assert.py mysql+pymysql://root:test@localhost/optunatest
optuna version: 2.3.0
...
RuntimeError: The runtime optuna version 2.3.0 is no longer compatible with the table schema (set up by optuna 2.3.0). Please execute `$ optuna storage upgrade --storage $STORAGE_URL` for upgrading the storage.

Upgrade the schema

$ optuna storage upgrade --storage mysql+pymysql://root:test@localhost/optunatest

Check the migration results

$ python migration-assert.py mysql+pymysql://root:test@localhost/optunatest
OK

Test for schema

Please run following two scripts. The first one execute migration-prepare.py with optuna@master while the second one execute it with optuna v2.3.0 and upgrade the DB.

docker rm -f some-mysql
docker run --name some-mysql -e MYSQL_ROOT_PASSWORD=test -p 3306:3306 -p 33060:33060 -d mysql:5.7
sleep 20
docker run --network host -it --rm mysql:5.7  mysql -h 127.0.0.1 -uroot -p -e "create database optunatest;"
pip install git+https://github.com/optuna/optuna.git
python migration-prepare.py mysql+pymysql://root:test@localhost/optunatest
mysqldump --skip-column-statistics -u root -p -h 127.0.0.1 optunatest > master.mysql

docker rm -f some-mysql
docker run --name some-mysql -e MYSQL_ROOT_PASSWORD=test -p 3306:3306 -p 33060:33060 -d mysql:5.7
sleep 20
docker run --network host -it --rm mysql:5.7  mysql -h 127.0.0.1 -uroot -p -e "create database optunatest;"
pip install optuna
python migration-prepare.py mysql+pymysql://root:test@localhost/optunatest
pip install git+https://github.com/optuna/optuna.git
optuna storage upgrade --storage mysql+pymysql://root:test@localhost/optunatest
echo "test" | mysqldump --skip-column-statistics -u root -p -h 127.0.0.1 optunatest > upgrade.mysql

The schema of these two should be identical except for trials' result. Please check the diff of the dump files, for example, as follows:

diff master.mysql upgrade.mysql | less -S

PostgreSQL

Test for values

Install the current version of optuna and create tables:

$ pip uninstall optuna  # if necessary
$ pip install optuna
$ docker run -it --rm --name postgres-test -e POSTGRES_PASSWORD=test -p 15432:5432 -d postgres
$ python migration-prepare.py postgresql+psycopg2://postgres:test@localhost:15432/postgres

Upgrade optuna to the patched version

$ pip uninstall optuna
$ pip install git+https://github.com/HideakiImamura/optuna.git@refactoring/change-rdb-schema-for-mo-refactoring

Check if RuntimeError occurs due to the incompatibility of the schema.

$ python migration-assert.py postgresql+psycopg2://postgres:test@localhost:15432/postgres
optuna version: 2.3.0
...
RuntimeError: The runtime optuna version 2.3.0 is no longer compatible with the table schema (set up by optuna 2.3.0). Please execute `$ optuna storage upgrade --storage $STORAGE_URL` for upgrading the storage.

Upgrade the schema

$ optuna storage upgrade --storage postgresql+psycopg2://postgres:test@localhost:15432/postgres

Check the migration results

$ python migration-assert.py postgresql+psycopg2://postgres:test@localhost:15432/postgres
OK

Test for schema

Please run following two scripts. The first one execute migration-prepare.py with optuna@master while the second one execute it with optuna v2.3.0 and upgrade the DB.

docker rm -f postgres-test
docker run -it --rm --name postgres-test -e POSTGRES_PASSWORD=test -p 15432:5432 -d postgres
sleep 10
pip install git+https://github.com/optuna/optuna.git
python migration-prepare.py postgresql+psycopg2://postgres:test@localhost:15432/postgres
pg_dump -h localhost -p 15432 -U postgres -d postgres > master.psql

docker rm -f postgres-test
docker run -it --rm --name postgres-test -e POSTGRES_PASSWORD=test -p 15432:5432 -d postgres
sleep 10
pip install optuna
python migration-prepare.py postgresql+psycopg2://postgres:test@localhost:15432/postgres
pip install git+https://github.com/optuna/optuna.git
optuna storage upgrade --storage postgres+psycopg2://postgres:test@localhost:15432/postgres
pg_dump -h localhost -p 15432 -U postgres -d postgres > upgrade.psql

The schema of these two should be identical except for trials' result. Please check the diff of the dump files, for example, as follows:

diff master.psql upgrade.psql | less -S

migration-assert.py

import sys

import optuna


if __name__ == "__main__":
    print(f"optuna version: {optuna.__version__}")
    study = optuna.load_study(study_name="test", storage=sys.argv[1])

    assert study.user_attrs["study_user_key"] == "study_user_value"
    assert study.system_attrs["study_system_key"] == "study_system_value"

    for i, trial in enumerate(study.trials):
        assert trial.user_attrs["user_key"] == "user_value"
        assert trial.system_attrs["system_key"] == "system_value"
        assert trial.number == i

        assert -100 <= trial.params["u"] < 100
        assert 1 <= trial.params["lu"] < 100
        assert -100 <= trial.params["du"] <= 100
        assert trial.params["du"] %2 == 0

        assert -100 <= trial.params["i"] <= 100
        assert 1 <= trial.params["li"] <= 100
        assert -100 <= trial.params["si"] <= 100
        assert trial.params["si"] %2 == 0

        assert trial.params["c"] in [-1, 0, 1]

        assert trial.intermediate_values == {i:2*i for i in range(10)}
    print("OK")

migration-prepare.py

import sys

import optuna


def objective(trial):
    u = trial.suggest_float("u", -100, 100)
    lu = trial.suggest_float("lu", 1, 100, log=True)
    du = trial.suggest_float("du", -100, 100, step=2)

    i = trial.suggest_int("i", -100, 100)
    li = trial.suggest_int("li", 1, 100, log=True)
    si = trial.suggest_int("si", -100, 100, step=2)

    c = trial.suggest_categorical("c", [-1, 0, 1])
    for epoch in range(10):
        trial.report(epoch * 2, epoch)

    trial.set_user_attr("user_key", "user_value")
    trial.set_system_attr("system_key", "system_value")

    return sum([u, lu, du, i, li, si, c])


if __name__ == "__main__":
    print(f"optuna version: {optuna.__version__}")
    study = optuna.create_study(study_name="test", storage=sys.argv[1])
    study.set_user_attr("study_user_key", "study_user_value")
    study.set_system_attr("study_system_key", "study_system_value")

    study.optimize(objective, n_trials=20)
    print(f"Best value: {study.best_value} (params: {study.best_params})\n")

sqlite_extract_create.py

import sys

tables = []

lines = ""
for line in open(sys.argv[1]):
    lines += line
    if line.endswith(";\n"):
        if lines.startswith("CREATE"):
            tables.append(lines)
        lines = ""

def sortkey(sql):
    # sql is a string such as
    # CREATE TABLE trials (
    # CREATE TABLE IF NOT EXISTS "trials" (
    if "(" in sql:
        table_name = sql.split("(")[0].split()[-1]
    else:
        table_name = sql.split()[2]
    if table_name.startswith('"'):
        return table_name[1:-1]
    return table_name

for table in sorted(tables, key=sortkey):
    print(table)