From f3d2ed8baaab5129e143ff8841a1c27c453b345b Mon Sep 17 00:00:00 2001 From: thomas Date: Sat, 17 Dec 2022 22:29:06 +0100 Subject: [PATCH 1/3] update --- src/lightning_app/core/work.py | 10 ++++++++++ src/lightning_app/runners/backends/mp_process.py | 6 ++++++ 2 files changed, 16 insertions(+) diff --git a/src/lightning_app/core/work.py b/src/lightning_app/core/work.py index 029f01fd2f7ae..ab4899061ccfb 100644 --- a/src/lightning_app/core/work.py +++ b/src/lightning_app/core/work.py @@ -618,6 +618,16 @@ def stop(self): app = _LightningAppRef().get_current() self._backend.stop_work(app, self) + def delete(self): + """Stops LightingWork component and shuts down hardware provisioned via L.CloudCompute.""" + if not self._backend: + raise Exception( + "Can't stop the work, it looks like it isn't attached to a LightningFlow. " + "Make sure to assign the Work to a flow instance." + ) + app = _LightningAppRef().get_current() + self._backend.delete_work(app, self) + def _check_run_is_implemented(self) -> None: if not is_overridden("run", instance=self, parent=LightningWork): raise TypeError( diff --git a/src/lightning_app/runners/backends/mp_process.py b/src/lightning_app/runners/backends/mp_process.py index dc0681390046e..36f3cb8097604 100644 --- a/src/lightning_app/runners/backends/mp_process.py +++ b/src/lightning_app/runners/backends/mp_process.py @@ -88,6 +88,9 @@ def stop_work(self, app, work: "lightning_app.LightningWork") -> None: work_manager: MultiProcessWorkManager = app.processes[work.name] work_manager.kill() + def delete_work(self, app, work: "lightning_app.LightningWork") -> None: + self.stop_work(app, work) + class CloudMultiProcessingBackend(MultiProcessingBackend): def __init__(self, *args, **kwargs): @@ -108,3 +111,6 @@ def stop_work(self, app, work: "lightning_app.LightningWork") -> None: disable_port(work._port) self.ports = [port for port in self.ports if port != work._port] return super().stop_work(app, work) + + def delete_work(self, app, work: "lightning_app.LightningWork") -> None: + self.stop_work(app, work) From 52fe778d289a62f57a3bdb2e079d51936512ee81 Mon Sep 17 00:00:00 2001 From: thomas Date: Sat, 17 Dec 2022 22:31:27 +0100 Subject: [PATCH 2/3] update --- src/lightning_app/CHANGELOG.md | 3 +++ src/lightning_app/core/work.py | 7 +++++-- 2 files changed, 8 insertions(+), 2 deletions(-) diff --git a/src/lightning_app/CHANGELOG.md b/src/lightning_app/CHANGELOG.md index b427988b92400..5e2bae87102fe 100644 --- a/src/lightning_app/CHANGELOG.md +++ b/src/lightning_app/CHANGELOG.md @@ -16,6 +16,9 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/). - Added more datatypes to serving component ([#16018](https://github.com/Lightning-AI/lightning/pull/16018)) +- Added `work.delete` method to delete the work ([#16103](https://github.com/Lightning-AI/lightning/pull/16103)) + + ### Changed - diff --git a/src/lightning_app/core/work.py b/src/lightning_app/core/work.py index ab4899061ccfb..d8d47af5a6e50 100644 --- a/src/lightning_app/core/work.py +++ b/src/lightning_app/core/work.py @@ -619,10 +619,13 @@ def stop(self): self._backend.stop_work(app, self) def delete(self): - """Stops LightingWork component and shuts down hardware provisioned via L.CloudCompute.""" + """Delete LightingWork component and shuts down hardware provisioned via L.CloudCompute. + + Locally, the work.delete() behaves as work.stop(). + """ if not self._backend: raise Exception( - "Can't stop the work, it looks like it isn't attached to a LightningFlow. " + "Can't delete the work, it looks like it isn't attached to a LightningFlow. " "Make sure to assign the Work to a flow instance." ) app = _LightningAppRef().get_current() From b2cf4fc26749c1ceff0517eb9852dfbeb76d975b Mon Sep 17 00:00:00 2001 From: thomas Date: Sat, 17 Dec 2022 22:35:59 +0100 Subject: [PATCH 3/3] update --- tests/tests_app/core/test_lightning_work.py | 14 +++++++++++++- 1 file changed, 13 insertions(+), 1 deletion(-) diff --git a/tests/tests_app/core/test_lightning_work.py b/tests/tests_app/core/test_lightning_work.py index cb97eabfa237c..695ef612d5125 100644 --- a/tests/tests_app/core/test_lightning_work.py +++ b/tests/tests_app/core/test_lightning_work.py @@ -1,6 +1,6 @@ from queue import Empty from re import escape -from unittest.mock import Mock +from unittest.mock import MagicMock, Mock import pytest @@ -384,3 +384,15 @@ def run(self): def test_lightning_app_work_start(cache_calls, parallel): app = LightningApp(FlowStart(cache_calls, parallel)) MultiProcessRuntime(app, start_server=False).dispatch() + + +def test_lightning_work_delete(): + work = WorkCounter() + + with pytest.raises(Exception, match="Can't delete the work"): + work.delete() + + mock = MagicMock() + work._backend = mock + work.delete() + assert work == mock.delete_work._mock_call_args_list[0].args[1]