Repository: Parsl/parsl

Issue:
`with` style context manager for Parsl DFK
**Is your feature request related to a problem? Please describe.**

Users should be able to use parsl DFKs using the `with` statement, to cause the DFK to be cleaned up when it is finished with.

Previously, if a user did not explicitly clean up a DFK, an `atexit` handler would do this when the process exits. See issue #3104 for why that doesn't work in recent Python, PR #3165 for the PR that removes that handling and issue #3251 for an example of a user encountering this as a practical problem.

I would like some syntax like this:

```
with parsl.load(config):
    f = my_app()
    f.result()
# implicit contextmanager cleanup of the loaded DFK happens here by calling dfk.cleanup()
```

In addition to introducing this context manager, documentation should be changed to encourage users to use this syntax instead of explicitly calling cleanup() in most scripted cases (although not always, for example, I think syntactically this doesn't work with jupyter notebooks as I think a `with` block cannot span multiple notebook cells)

**Describe alternatives you've considered**
It might be possible to rework exit handling in an atexit handler to still be able to do shutdown, but I don't feel like that will be easy - because a lot of stuff can happen at DFK shutdown.

**Additional context**
This is coming from changed (more strict) behaviour around threads and processes in Python 3.12.


Discussion:
@benclifford I'd like to work on this issue. Thank you.

@benclifford I will work on this issue if thats ok
ok if multiple people are working on this, you might all like to comment here on your approaches and try to work together a bit
> ok if multiple people are working on this, you might all like to comment here on your approaches and try to work together a bit

ok @benclifford, i just cloned the repo and trying to better understand the issue, will give details of my approach asap


Resolving patch:
diff --git a/parsl/dataflow/dflow.py b/parsl/dataflow/dflow.py
--- a/parsl/dataflow/dflow.py
+++ b/parsl/dataflow/dflow.py
@@ -206,6 +206,13 @@ def __init__(self, config: Config) -> None:
 
         atexit.register(self.atexit_cleanup)
 
+    def __enter__(self):
+        pass
+
+    def __exit__(self, exc_type, exc_value, traceback):
+        logger.debug("Exiting the context manager, calling cleanup for DFK")
+        self.cleanup()
+
     def _send_task_log_info(self, task_record: TaskRecord) -> None:
         if self.monitoring:
             task_log_info = self._create_task_log_info(task_record)


Test patch:
diff --git a/parsl/tests/test_python_apps/test_context_manager.py b/parsl/tests/test_python_apps/test_context_manager.py
new file mode 100644
--- /dev/null
+++ b/parsl/tests/test_python_apps/test_context_manager.py
@@ -0,0 +1,40 @@
+import parsl
+from parsl.tests.configs.local_threads import fresh_config
+import pytest
+from parsl.errors import NoDataFlowKernelError
+
+
+@parsl.python_app
+def square(x):
+    return x * x
+
+
+@parsl.bash_app
+def foo(x, stdout='foo.stdout'):
+    return f"echo {x + 1}"
+
+
+def local_setup():
+    pass
+
+
+def local_teardown():
+    parsl.clear()
+
+
+@pytest.mark.local
+def test_within_context_manger():
+    config = fresh_config()
+    with parsl.load(config=config):
+        py_future = square(2)
+        assert py_future.result() == 4
+
+        bash_future = foo(1)
+        assert bash_future.result() == 0
+
+        with open('foo.stdout', 'r') as f:
+            assert f.read() == "2\n"
+
+    with pytest.raises(NoDataFlowKernelError) as excinfo:
+        square(2).result()
+    assert str(excinfo.value) == "Cannot submit to a DFK that has been cleaned up"