Fix bug of custom user widgets and enhance artifact support.

[c-bata]

Contributor License Agreement

This repository (optuna-dashboard) and Goptuna share common code.
This pull request may therefore be ported to Goptuna.
Make sure that you understand the consequences concerning licenses and check the box below if you accept the term before creating this pull request.

• I agree this patch may be ported to Goptuna by other Goptuna contributors.

Reference Issues/PRs

What does this implement/fix? Explain your changes.

While I created a demo app for human-in-the-loop optimization, I found some issues. This PR fixes them.

from __future__ import annotations

import os
import textwrap
import threading
import time
from typing import NoReturn
from wsgiref.simple_server import make_server

import optuna
from PIL import Image
from optuna.trial import TrialState

from optuna_dashboard import ObjectiveChoiceWidget, save_note, ObjectiveSliderWidget
from optuna_dashboard import register_objective_form_widgets
from optuna_dashboard import set_objective_names
from optuna_dashboard import wsgi
from optuna_dashboard.artifact import upload_artifact
from optuna_dashboard.artifact.file_system import FileSystemBackend


storage = optuna.storages.InMemoryStorage()
base_path = os.path.join(os.path.dirname(__file__), "artifact")
artifact_backend = FileSystemBackend(base_path=base_path)


def suggest_and_generate_image(study: optuna.Study) -> None:
    # Ask new parameters
    trial = study.ask()
    r = trial.suggest_int("r", 0, 255)
    g = trial.suggest_int("g", 0, 255)
    b = trial.suggest_int("b", 0, 255)

    # Generate image
    image_path = f"tmp/sample-{trial.number}.png"
    image = Image.new("RGB", (320, 240), color=(r,g,b))
    image.save(image_path)

    # Upload Artifact
    artifact_id = upload_artifact(artifact_backend, trial, image_path)

    # Save Note
    note = textwrap.dedent(f'''\
    ## Trial {trial.number}
    
    ![generated-image](/artifacts/{study._study_id}/{trial._trial_id}/{artifact_id})
    ''')
    save_note(trial, note)


def start_preferential_optimization() -> NoReturn:
    # Create Study
    study = optuna.create_study(
        study_name="Preferential Optimization",
        storage=storage,
        directions=["minimize", "maximize"]
    )
    set_objective_names(study, ["Human Perception Score", "Looks Yellow?"])
    register_objective_form_widgets(study, widgets=[
        ObjectiveChoiceWidget(
            choices=["Good 👍", "Bad 👎"],
            values=[-1, 1],
            description="Please input your score!",
        ),
        ObjectiveSliderWidget(
            min=1,
            max=10,
            step=1,
            description="Score 10 if it looks YELLOW!",
        ),
    ])

    # Start Preferential Optimization
    n_batch = 2
    while True:
        running_trials = study.get_trials(deepcopy=False, states=(TrialState.RUNNING,))
        if len(running_trials) >= n_batch:
            time.sleep(10)
            continue
        suggest_and_generate_image(study)


def main() -> None:
    if not os.path.exists(base_path):
        os.mkdir(base_path)

    # Start Dashboard server on background
    app = wsgi(storage, artifact_backend=artifact_backend)
    httpd = make_server("127.0.0.1", 8080, app)
    thread = threading.Thread(target=httpd.serve_forever)
    thread.start()

    # Run optimize loop
    try:
        start_preferential_optimization()
    except KeyboardInterrupt:
        httpd.shutdown()
        httpd.server_close()
        thread.join()


if __name__ == "__main__":
    main()