c-bata/profile_optuna_pr_4449_on_postgresql.py

## profile_optuna_pr_4449_on_postgresql.py
# $ docker run -d --rm --platform linux/amd64 -p 5432:5432 -e POSTGRES_USER=root -e POSTGRES_PASSWORD=root -e POSTGRES_DB=optuna --name optuna-postgres postgres:12.10

# $ docker run -it --rm --platform linux/amd64 --network host -v $(pwd):/usr/src python:3.10 bash
# # cd /usr/src
# # pip install -U setuptools pip psycopg2
# # pip install -e .

from __future__ import annotations

import math
import threading
import time

from sqlalchemy import event
from sqlalchemy.engine.base import Engine

import optuna


optuna.logging.set_verbosity(optuna.logging.ERROR)

storage_url = "postgresql+psycopg2://root:root@127.0.0.1/optuna"
storage = optuna.storages.RDBStorage(storage_url)

sql_queries_lock = threading.Lock()
sql_queries: dict[str, tuple[int, list[float]]] = {}

n_studies = 100
n_trials = 500
n_params = 10


class EngineProfiler:
    def __init__(self, engine: Engine) -> None:
        self.engine = engine
        self.query_start_time = time.perf_counter()

    def register(self) -> None:
        event.listen(self.engine, "before_cursor_execute", self.before_cursor_execute)
        event.listen(self.engine, "after_cursor_execute", self.after_cursor_execute)

    def before_cursor_execute(  # type: ignore
        self, conn, cursor, statement, parameters, context, executemany
    ) -> None:
        self.query_start_time = time.perf_counter()

    def after_cursor_execute(  # type: ignore
            self, conn, cursor, stmt, parameters, context, executemany
    ) -> None:
        global sql_queries, sql_queries_lock

        duration = time.perf_counter() - self.query_start_time
        with sql_queries_lock:
            registered = stmt in sql_queries
            sql_queries[stmt] = (
                sql_queries[stmt][0] + 1 if registered else 1,
                sql_queries[stmt][1] + [duration] if registered else [duration],
            )


def objective(trial: optuna.Trial) -> float:
    return sum([
        math.sin(trial.suggest_float('param-{}'.format(i), 0, math.pi * 2))
        for i in range(n_params)
    ])


def main():
    global sql_queries, sql_queries_lock
    # Create trials
    if len(storage.get_all_studies()) == 0:
        for i in range(n_studies):
            print(i)
            study = optuna.create_study(storage=storage)
            study.optimize(lambda trial: objective(trial), n_trials=n_trials, n_jobs=8)

    # Profile study.optimize()
    EngineProfiler(storage.engine).register()
    start = time.time()
    tmp_studies = []
    for i in range(100):
        tmp_study = optuna.create_study(storage=storage)
        tmp_study.optimize(objective, n_trials=10)
        tmp_studies.append(tmp_study)
    elapsed = time.time() - start
    print(f"Elapsed: {elapsed:.4f}s ({n_trials=} {n_params=})")

    # Show profiler stats
    summary = [
        (stmt, count, f"{sum(durations):.4f}", sum(durations))
        for stmt, (count, durations) in sql_queries.items()
    ]
    sort_by_total = sorted(summary, key=lambda r: r[3], reverse=True)
    print("Sort by Total:")
    print("Total Time(s)\tQuery Count\tStatement")
    for q in sort_by_total[:5]:
        print(f"{q[2]}\t{q[1]}\t{q[0]}")

    # Clean up
    for s in tmp_studies:
        optuna.delete_study(study_name=s.study_name, storage=storage)


if __name__ == '__main__':
    main()
	# $ docker run -d --rm --platform linux/amd64 -p 5432:5432 -e POSTGRES_USER=root -e POSTGRES_PASSWORD=root -e POSTGRES_DB=optuna --name optuna-postgres postgres:12.10

	# $ docker run -it --rm --platform linux/amd64 --network host -v $(pwd):/usr/src python:3.10 bash
	# # cd /usr/src
	# # pip install -U setuptools pip psycopg2
	# # pip install -e .

	from __future__ import annotations

	import math
	import threading
	import time

	from sqlalchemy import event
	from sqlalchemy.engine.base import Engine

	import optuna


	optuna.logging.set_verbosity(optuna.logging.ERROR)

	storage_url = "postgresql+psycopg2://root:root@127.0.0.1/optuna"
	storage = optuna.storages.RDBStorage(storage_url)

	sql_queries_lock = threading.Lock()
	sql_queries: dict[str, tuple[int, list[float]]] = {}

	n_studies = 100
	n_trials = 500
	n_params = 10


	class EngineProfiler:
	def __init__(self, engine: Engine) -> None:
	self.engine = engine
	self.query_start_time = time.perf_counter()

	def register(self) -> None:
	event.listen(self.engine, "before_cursor_execute", self.before_cursor_execute)
	event.listen(self.engine, "after_cursor_execute", self.after_cursor_execute)

	def before_cursor_execute( # type: ignore
	self, conn, cursor, statement, parameters, context, executemany
	) -> None:
	self.query_start_time = time.perf_counter()

	def after_cursor_execute( # type: ignore
	self, conn, cursor, stmt, parameters, context, executemany
	) -> None:
	global sql_queries, sql_queries_lock

	duration = time.perf_counter() - self.query_start_time
	with sql_queries_lock:
	registered = stmt in sql_queries
	sql_queries[stmt] = (
	sql_queries[stmt][0] + 1 if registered else 1,
	sql_queries[stmt][1] + [duration] if registered else [duration],
	)


	def objective(trial: optuna.Trial) -> float:
	return sum([
	math.sin(trial.suggest_float('param-{}'.format(i), 0, math.pi * 2))
	for i in range(n_params)
	])


	def main():
	global sql_queries, sql_queries_lock
	# Create trials
	if len(storage.get_all_studies()) == 0:
	for i in range(n_studies):
	print(i)
	study = optuna.create_study(storage=storage)
	study.optimize(lambda trial: objective(trial), n_trials=n_trials, n_jobs=8)

	# Profile study.optimize()
	EngineProfiler(storage.engine).register()
	start = time.time()
	tmp_studies = []
	for i in range(100):
	tmp_study = optuna.create_study(storage=storage)
	tmp_study.optimize(objective, n_trials=10)
	tmp_studies.append(tmp_study)
	elapsed = time.time() - start
	print(f"Elapsed: {elapsed:.4f}s ({n_trials=} {n_params=})")

	# Show profiler stats
	summary = [
	(stmt, count, f"{sum(durations):.4f}", sum(durations))
	for stmt, (count, durations) in sql_queries.items()
	]
	sort_by_total = sorted(summary, key=lambda r: r[3], reverse=True)
	print("Sort by Total:")
	print("Total Time(s)\tQuery Count\tStatement")
	for q in sort_by_total[:5]:
	print(f"{q[2]}\t{q[1]}\t{q[0]}")

	# Clean up
	for s in tmp_studies:
	optuna.delete_study(study_name=s.study_name, storage=storage)


	if __name__ == '__main__':
	main()