Executing code before rerun

[Amadan]

I am launching a long celery task, and displaying the results when it is done. If the user changes input (and thus reruns the streamlit app), I wish to detect the rerun, and cancel the launched task before the rerun is initiated (using a variable defined in the run being canceled).

I was reading the code to find out how to do it, but I couldn’t quite manage to figure it out. I thought maybe I could connect to script_runner.on_event and listen for ScriptRunnerEvent. SCRIPT_STOPPED_FOR_RERUN, but I couldn’t find out how to get access to script_runner from my script (or if it is even possible). I would be grateful if someone has a solution for this.

As a workaround, I tried to put my celery task id into st.session_state, but that failed too: it seems that st.session_state is not persisted when the script has been rerun. Here is a mock-up of that situation:

import streamlit as st
from random import randrange
from time import sleep

task_id = st.session_state.get("n")
if task_id is not None:
    print(f"Cancelling previous task {task_id}")

def process(input_data):
    task_id = randrange(256)
    st.markdown(f"`{input_data}` being processed by task `{task_id}`")
    st.session_state.task_id = task_id
    bar = st.progress(0)
    for t in range(100):
        sleep(0.1)
        bar.progress(t + 1)
    del st.session_state.task_id
    st.markdown(f"`{input_data}` processed successfully.")

input_data = st.text_input("Input Data")
if input_data:
    process(input_data)

If new input data is entered before the progress bar hits 100, the application reruns, but last_id is None. Is there a way to “flush” the session state after st.session_state.task_id = task_id, to force it to persist even if a rerun is triggered later?

[Amadan]

Okay, I think I figured it out.

from streamlit.runtime.scriptrunner import RerunException, StopException

try:
    launch_task()
except (RerunException, StopException):
    cleanup_task()
    raise