diff --git a/python/examples/executor.py b/python/examples/executor.py index f926ee8e9..bb032f634 100644 --- a/python/examples/executor.py +++ b/python/examples/executor.py @@ -27,3 +27,5 @@ executor.wait_for_termination() # %% # %% +executor.close() +# %% diff --git a/python/examples/scheduler.py b/python/examples/scheduler.py index df3bf150e..1c40ce1ee 100644 --- a/python/examples/scheduler.py +++ b/python/examples/scheduler.py @@ -26,3 +26,4 @@ # %% scheduler.wait_for_termination() # %% +scheduler.close() \ No newline at end of file diff --git a/python/src/cluster.rs b/python/src/cluster.rs index b2b065629..e395e45f6 100644 --- a/python/src/cluster.rs +++ b/python/src/cluster.rs @@ -96,6 +96,19 @@ impl PyScheduler { None => Ok(()), } } + + pub fn close(&mut self) -> PyResult<()> { + let mut handle = None; + std::mem::swap(&mut self.handle, &mut handle); + + match handle { + Some(handle) => handle.abort(), + None => {} + }; + + Ok(()) + } + #[classattr] pub fn version() -> &'static str { ballista_core::BALLISTA_VERSION @@ -199,6 +212,18 @@ impl PyExecutor { } } + pub fn close(&mut self) -> PyResult<()> { + let mut handle = None; + std::mem::swap(&mut self.handle, &mut handle); + + match handle { + Some(handle) => handle.abort(), + None => {} + }; + + Ok(()) + } + #[classattr] pub fn version() -> &'static str { ballista_core::BALLISTA_VERSION