Add streamlit helper

[cross32768]

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.

What does this implement/fix? Explain your changes.

This PR adds streamlit helper for optuna-dashboard.
Using this helper, users can render some widgets on streamlit with simple interface as the image below.

I added 3 streamlit helper functions and its unit tests.

• render_trial_note
  • This function write a trial note to UI with streamlit as a markdown format. It is helpful to use with save_note function already in this repo (see: test_render_trial_note).
• render_user_attr_form_widgets
  • This function render user input widgets to UI with streamlit. Submitted values to the forms are registered as each trial's user_attrs. Is is helpful to use with register_user_attr_form_widgets function already in this repo (see: test_render_user_attr_form_widgets).
• render_objective_form_widgets
  • This function render user input widgets to UI with streamlit. Submitted values to the forms are telled to optuna trial object. It is helpful to use with register_objective_form_widgets function already in this repo (see: test_render_objective_form_widgets)

[c-bata]

Thank you for your pull request!
I understand this PR is still WIP, but I left an early feedback comment.

[c-bata]

Could you add the optuna.streamlit package and put these functions in it? Then we can make streamlit as an optional dependency.

[cross32768]

OK. Does you mean I create directory "optuna-dashboard/streamlit" and put my code in this dir instead of "optuna-dashboard/optuna-dashboard/_streamlit_helper.py"?

[c-bata]

Exactly 👍

[cross32768]

You can test the work of this PR with the code below (This code is to test render_trial_note and render_user_attr_form_widgets).

human_in_the_worker.py

import time

import optuna
from optuna.trial import TrialState
from optuna_dashboard import ChoiceWidget
from optuna_dashboard import ObjectiveUserAttrRef
from optuna_dashboard import register_user_attr_form_widgets
from optuna_dashboard import save_note
from optuna_dashboard import SliderWidget
from optuna_dashboard import TextInputWidget


def main():
    study = optuna.create_study(study_name="test", storage="sqlite:///example.db")

    widgets = [
        ChoiceWidget(
            choices=["Good", "Bad"],
            values=[1, -1],
            description="description",
            user_attr_key="choice",
        ),
        SliderWidget(
            min=1,
            max=5,
            step=1,
            labels=[(1, "Bad"), (5, "Good")],
            description="description",
            user_attr_key="slider",
        ),
        TextInputWidget(description="description", user_attr_key="text1"),
        TextInputWidget(description="description", user_attr_key="text2"),
    ]
    register_user_attr_form_widgets(
        study,
        widgets=widgets,
    )

    n_batch = 5
    while True:
        # Check trials
        running_trials = study.get_trials(deepcopy=False, states=(TrialState.RUNNING,))
        for t in running_trials:
            if "choice" in t.user_attrs:
                study.tell(t.number, 1)

        # Create a new trial
        running_trials = study.get_trials(deepcopy=False, states=(TrialState.RUNNING,))
        if len(running_trials) < n_batch:
            trial = study.ask()
            save_note(trial, f"test message for {trial.number}")

        # Avoid busy-loop
        time.sleep(5)


if __name__ == "__main__":
    main()

streamlit_worker.py

import warnings

import optuna
from optuna.trial import TrialState
from optuna_dashboard.streamlit import render_trial_note
from optuna_dashboard.streamlit import render_user_attr_form_widgets

import streamlit as st


warnings.simplefilter("ignore", FutureWarning)


def start_streamlit() -> None:
    study = optuna.load_study(storage="sqlite:///example.db", study_name="test")
    selected_trial = st.sidebar.selectbox(
        "trial number",
        study.trials,
        format_func=lambda t: t.number,
    )

    render_trial_note(study, selected_trial)

    if selected_trial.state == TrialState.RUNNING:
        render_user_attr_form_widgets(study, selected_trial)
    if selected_trial.state == TrialState.COMPLETE:
        st.write(selected_trial.user_attrs)


if __name__ == "__main__":
    start_streamlit()

Please run python3 human_in_the_worker.py and streamlit run streamlit_worker.py

[c-bata]

@knshnb Could you review this PR?

[c-bata]

LGTM!

I created a simple application based on human-in-the-loop tutorial and confirmed that this PR works as expected.

https://gist.github.com/c-bata/ffcf52ca5c6fee5ec64e56eee451aabd