Streamlit如何实现有状态ML应用 避免组件交互时全脚本重跑

核心原理 #

Streamlit默认每次交互都会全量重跑整个脚本，st.button的返回值仅在点击的那次运行中为True，后续交互会重置为False，同时你定义的df变量会被重置为空，导致之前的加载、计算逻辑全部重新执行。
我们可以通过内置的st.session_state做状态持久化 + 函数缓存来解决这个问题，仅让交互关联的直方图模块重跑。

修改后的完整代码 #

import numpy as np
import pandas as pd
import sklearn as sk
import matplotlib.pyplot as plt
import streamlit as st
import pyspark
from pyspark import *
from PIL import Image
from io import StringIO

# 加缓存装饰器，相同参数下仅执行一次，不可哈希的组件参数加下划线跳过哈希校验
@st.cache_resource
def load_data(ss,uploaded_file):
    df = ss.read.format('csv').option('header','true').load(uploaded_file)
    return df

@st.cache_data
def sample_data(df,_widget):
    df_sample = pd.DataFrame(df.head(5))
    df_sample.columns = df.columns
    _widget.dataframe(df_sample)
    
@st.cache_data
def descriptive_stats(df,_widget):
    df_desc = df.summary().toPandas()
    _widget.dataframe(df_desc)
    
def hist_plot(df,col,_widget):
    df_plot = df.select(col).toPandas().iloc[:,0]
    fig, ax = plt.subplots()
    ax.hist(df_plot,density = False, bins = 50)
    _widget.pyplot(fig)

def main():
    # 初始化会话状态，持久化存储已加载数据和上传完成标记
    if "df" not in st.session_state:
        st.session_state.df = None
    if "upload_done" not in st.session_state:
        st.session_state.upload_done = False

    sparkapp = pyspark.sql.SparkSession.builder.master('local[4]').appName('No-code Spark Pipeline').getOrCreate()
    st.title("No-Code ML Spark Pipeline")
    st.subheader('1. Upload file (csv)')
    uploaded_file = st.text_input("Provide local file path")
    upload_button1 = st.button('Upload')
    st.caption('Sample data')
    upload_cont1 = st.empty()
    # 可替换为你本地的占位图片
    # white_background = Image.open('C:/Users/hp/Desktop/white_600_240.png')
    # upload_cont1.image(white_background)

    # 点击上传按钮时仅执行一次加载逻辑，结果存入session_state
    if upload_button1:
        df = load_data(sparkapp,uploaded_file)
        st.session_state.df = df
        st.session_state.upload_done = True

    # 上传完成后常驻展示样本，不会重复计算
    if st.session_state.upload_done:
        sample_data(st.session_state.df, upload_cont1)
    
    st.subheader('2. Exploratory Data Analytics')
    st.caption('Descriptive statistics')
    eda_cont1 = st.empty()
    # eda_cont1.image(white_background,use_column_width=True)
    
    # 上传完成后常驻展示统计结果，不会重复计算
    if st.session_state.upload_done:
        descriptive_stats(st.session_state.df, eda_cont1)         
    
    st.caption('Histogram / Frequency plot')
    if st.session_state.upload_done:
        eda_sel_feat = st.selectbox('Select feature to be displayed', options = st.session_state.df.columns)
        eda_cont2 = st.empty()
        # 仅切换特征时重跑绘图逻辑
        hist_plot(st.session_state.df,eda_sel_feat,eda_cont2)


if __name__ == '__main__':
    main()

改动说明 #

• 新增st.session_state存储上传后的数据集和完成状态，后续任意交互都不会丢失已加载的数据
• 给数据加载、样本生成、描述性统计函数加缓存装饰器，相同参数下仅会执行一次，不会重复计算
• 样本、统计结果的展示逻辑调整为上传完成后常驻，不会因为下拉框交互消失
• 直方图仅和下拉框选择的特征关联，每次切换特征仅重跑绘图逻辑，前面的步骤完全不会重复执行

内容的提问来源于stack exchange，提问作者bhargav