Step 4: Implement the UI with StreamLit¶
In step 3, we have implemented a Flow which dynamically creates a Work when a new request is added to the requests list.
From the UI, we create 3 pages with StreamLit:
Page 1: Create a form with add a new request to the Flow state requests.
Page 2: Iterate through all the requests and display the associated information.
Page 3: Display the entire App State.
Render All Pages¶
def render_fn(state: AppState):
import streamlit as st
page_names_to_funcs = {
"Create a new Run": partial(page_1__create_new_run, state=state),
"View your Runs": partial(page_2__view_run_lists, state=state),
"View the App state": partial(page_3__view_app_state, state=state),
}
selected_page = st.sidebar.selectbox(
"Select a page", page_names_to_funcs.keys())
page_names_to_funcs[selected_page]()
Page 1
def page_1__create_new_run(state):
import streamlit as st
st.markdown("# Create a new Run 🎈")
# 1: Collect arguments from the users
id = st.text_input("Name your run", value="my_first_run")
github_repo = st.text_input(
"Enter a Github Repo URL", value="https://github.com/Lightning-AI/lightning-quick-start.git"
)
default_script_args = (
"--trainer.max_epochs=5"
" --trainer.limit_train_batches=4"
" --trainer.limit_val_batches=4"
" --trainer.callbacks=ModelCheckpoint"
" --trainer.callbacks.monitor=val_acc"
)
default_requirements = "torchvision, pytorch_lightning, jsonargparse[signatures]"
script_path = st.text_input("Enter your script to run", value="train_script.py")
script_args = st.text_input("Enter your base script arguments", value=default_script_args)
requirements = st.text_input("Enter your requirements", value=default_requirements)
ml_framework = st.radio(
"Select your ML Training Frameworks", options=["PyTorch Lightning", "Keras", "Tensorflow"]
)
if ml_framework not in ("PyTorch Lightning"):
st.write(f"{ml_framework} isn't supported yet.")
return
clicked = st.button("Submit")
# 2: If clicked, create a new request.
if clicked:
new_request = {
"id": id,
"train": {
"github_repo": github_repo,
"script_path": script_path,
"script_args": script_args.split(" "),
"requirements": requirements.split(" "),
"ml_framework": ml_framework,
},
}
# 3: IMPORTANT: Add a new request to the state in-place.
# The flow receives the UI request and dynamically create
# and run the associated work from the request information.
state.requests = state.requests + [new_request]
Page 2
def page_2__view_run_lists(state):
import streamlit as st
st.markdown("# Run Lists 🎈")
# 1: Iterate through all the requests in the state.
for i, r in enumerate(state.requests):
i = str(i)
# 2: Display information such as request, logs, work state, model score.
work = state._state["structures"]["ws"]["works"][f"w_{i}"]
with st.expander(f"Expand to view Run {i}", expanded=False):
if st.checkbox("Expand to view your configuration", key=i):
st.json(r)
if st.checkbox("Expand to view logs", key=i):
st.code(body=work["vars"]["logs"])
if st.checkbox("Expand to view your work state", key=i):
work["vars"].pop("logs")
st.json(work)
best_model_score = r.get("best_model_score", None)
if best_model_score:
if st.checkbox("Expand to view your run performance", key=i):
st.json({"best_model_score": best_model_score, "best_model_path": r.get("best_model_path")})
Page 3
def page_3__view_app_state(state):
import streamlit as st
st.markdown("# App State 🎈")
st.write(state._state)
