注册的atexit函数无法访问Streamlit Session State变量问题

问题描述 #

我通过如下方式注册了一个atexit函数：

import streamlit as st
from azure.search.documents.indexes import SearchIndexClient 
from azure.core.credentials import AzureKeyCredential


@atexit.register
def delete_vector_index():
    admin_client = SearchIndexClient(endpoint=st.session_state.vector_store_address,
                      index_name=st.session_state.index_name,
                      credential=AzureKeyCredential(st.session_state.vector_store_password))
    admin_client.delete_index(st.session_state.index_name)

if "initialized" not in st.session_state:
    st.session_state.initialized = True
    st.session_state.vector_store_address: str = "<my endpoint>"
    st.session_state.vector_store_password: str = "<admin key>"
    st.session_state.index_name: str = "<index name>"

但每次解释器运行结束时，都会出现报错：AttributeError: st.session_state has no attribute "vector_store_address". Did you forget to initialize it?

为何该atexit函数无法访问st.session_state变量，而我的其他函数却可以正常访问？请问该如何解决这个问题？

原因分析 #

这是因为Python的atexit钩子执行时机晚于Streamlit的会话上下文销毁时机：

• 业务函数运行时处于Streamlit的会话生命周期内，此时st.session_state是有效的；
• 当程序退出触发atexit函数时，Streamlit已经清理了会话相关资源，st.session_state已被重置或销毁，自然无法读取其中的属性。

解决方法 #

方法1：提前将参数保存到全局变量 #

在初始化st.session_state的同时，把需要的参数同步到全局变量，让atexit函数直接读取全局变量：

import atexit
import streamlit as st
from azure.search.documents.indexes import SearchIndexClient 
from azure.core.credentials import AzureKeyCredential

# 定义全局变量存储参数
vector_store_address = ""
vector_store_password = ""
index_name = ""

@atexit.register
def delete_vector_index():
    admin_client = SearchIndexClient(endpoint=vector_store_address,
                      index_name=index_name,
                      credential=AzureKeyCredential(vector_store_password))
    admin_client.delete_index(index_name)

if "initialized" not in st.session_state:
    st.session_state.initialized = True
    st.session_state.vector_store_address = "<my endpoint>"
    st.session_state.vector_store_password = "<admin key>"
    st.session_state.index_name = "<index name>"
    # 同步参数到全局变量
    global vector_store_address, vector_store_password, index_name
    vector_store_address = st.session_state.vector_store_address
    vector_store_password = st.session_state.vector_store_password
    index_name = st.session_state.index_name

方法2：使用Streamlit会话结束回调（推荐） #

Streamlit 1.28及以上版本提供了st.session_state.on_session_end回调，该回调会在会话正常结束时触发，此时st.session_state仍然有效：

import streamlit as st
from azure.search.documents.indexes import SearchIndexClient 
from azure.core.credentials import AzureKeyCredential

def delete_vector_index():
    admin_client = SearchIndexClient(endpoint=st.session_state.vector_store_address,
                      index_name=st.session_state.index_name,
                      credential=AzureKeyCredential(st.session_state.vector_store_password))
    admin_client.delete_index(st.session_state.index_name)

if "initialized" not in st.session_state:
    st.session_state.initialized = True
    st.session_state.vector_store_address = "<my endpoint>"
    st.session_state.vector_store_password = "<admin key>"
    st.session_state.index_name = "<index name>"
    # 注册会话结束回调
    st.session_state.on_session_end(delete_vector_index)

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