Clicking "Stop Pipeline" should terminate the code_map operation gracefully

[plpycoin]

• After starting a code_map operation, clicking "Stop Pipeline" during the processing of the operation will not stop the code_map thread.

[reproduce-bot]

The following script is generated by AI Agent to help reproduce the issue:

# docetl/reproduce.py
import threading
import time
from concurrent.futures import ThreadPoolExecutor

class MockPipeline:
    def __init__(self):
        self.stop_event = threading.Event()
        self.executor = ThreadPoolExecutor(max_workers=2)

    def code_map_operation(self):
        def mock_task():
            while not self.stop_event.is_set():
                time.sleep(0.1)  # Simulate work
            return "Task Stopped"

        futures = [self.executor.submit(mock_task) for _ in range(2)]
        return futures

    def check_if_tasks_running(self, futures):
        for future in futures:
            if not future.done():
                return True
        return False

    def stop_pipeline(self):
        self.stop_event.set()
        self.executor.shutdown(wait=False)

def test_stop_pipeline():
    pipeline = MockPipeline()
    futures = pipeline.code_map_operation()

    time.sleep(1)  # Let the tasks run for a while
    pipeline.stop_pipeline()  # Attempt to stop the pipeline

    tasks_running = pipeline.check_if_tasks_running(futures)
    if tasks_running:
        raise AssertionError("Tasks are still running after attempting to stop the pipeline")

    try:
        for future in futures:
            result = future.result(timeout=1)  # This should raise an exception if the tasks don't stop
            assert result == "Task Stopped"
        print("Test passed successfully with no errors!")
    except Exception as e:
        raise AssertionError(e)

if __name__ == "__main__":
    test_stop_pipeline()

How to run:

python3 docetl/reproduce.py

Expected Result:

Traceback (most recent call last):
  File "docetl/reproduce.py", line 49, in <module>
    test_stop_pipeline()
  File "docetl/reproduce.py", line 38, in test_stop_pipeline
    raise AssertionError("Tasks are still running after attempting to stop the pipeline")
AssertionError: Tasks are still running after attempting to stop the pipeline

Thank you for your valuable contribution to this project and we appreciate your feedback! Please respond with an emoji if you find this script helpful. Feel free to comment below if any improvements are needed.

Best regards from an AI Agent!
@plpycoin