Handling button events while receiving data from a background thread

I’m new to using Streamlit and am hoping that it can work for my use-case. I have a background thread that receives data from a process that I’m monitoring. By using a loop I’m able to periodically update streamlit with updated state information. I also have a button to interact with the background process.

The problem I face is that for the button’s press event to be captured, it needs to be inside the update loop. However, once the button is pressed it remains in the pressed state until the script returns. The work around has been to have the script return when the button is pressed and include a “restart” button that will re-run the script. I’ve included an example below that re-creates this flow.

Is there a cleaner way to make this work? I.e. is there a way to tell Streamlit to re-run the script programatically?

import streamlit as st
import time
import threading


class DataGenerator:
    def __init__(self):
        self.counter = 0

    def run(self):
        thread = threading.Thread(target=self._run, args=(), daemon=True)
        thread.start()

    def _run(self):
        while True:
            time.sleep(0.01)
            self.counter += 1

    def run_task(self):
        print(f"Do the thing once {self.counter}")


@st.cache(allow_output_mutation=True)
def generate_data():
    gen = DataGenerator()
    gen.run()
    return gen

# Declare UI Elements
running_time = st.empty()
generated_data = st.empty()
task_button = st.button("Do the thing")
kick_button = st.button("Restart")

# Start data generation
my_data = generate_data()

# Continually update the UI as data arrives from the background thread
while True:
    time.sleep(0.1)

    running_time.write(time.time())
    generated_data.write(my_data.counter)

    if task_button:
        my_data.run_task()
        break