addition to App 0.7 (#15219)

[App] Pass LightningWork to LightningApp (#15215) * update * update * update * Apply suggestions from code review * Apply suggestions from code review * Apply suggestions from code review * ll Co-authored-by: Jirka Borovec <[email protected]> Co-authored-by: Jirka <[email protected]> Co-authored-by: Luca Antiga <[email protected]> (cherry picked from commit 8ec7ddf) Co-authored-by: thomas chaton <[email protected]>
Lightning-AI · Oct 20, 2022 · 65d29f0 · 65d29f0
1 parent b73087c
commit 65d29f0
Show file tree

Hide file tree

Showing 4 changed files with 39 additions and 4 deletions.
diff --git a/src/lightning_app/CHANGELOG.md b/src/lightning_app/CHANGELOG.md
@@ -4,13 +4,14 @@ All notable changes to this project will be documented in this file.
 
 The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
 
-## [0.7.0] - 2022-10-19
+## [0.7.0] - 2022-10-20
 
 ### Added
 
 - Add `--secret` option to CLI to allow binding Secrets to app environment variables when running in the cloud ([#14612](https://github.com/Lightning-AI/lightning/pull/14612))
 - Added support for adding descriptions to commands either through a docstring or the `DESCRIPTION` attribute ([#15193](https://github.com/Lightning-AI/lightning/pull/15193)
 - Added option to add custom meta tags to the UI container ([#14915](https://github.com/Lightning-AI/lightning/pull/14915))
+- Added support to pass a `LightningWork` to the `LightningApp` ([#15215](https://github.com/Lightning-AI/lightning/pull/15215)
 
 ### Changed
 

diff --git a/src/lightning_app/core/app.py b/src/lightning_app/core/app.py
@@ -46,7 +46,7 @@
 class LightningApp:
     def __init__(
         self,
-        root: "lightning_app.LightningFlow",
+        root: "t.Union[lightning_app.LightningFlow, lightning_app.LightningWork]",
         debug: bool = False,
         info: frontend.AppInfo = None,
         root_path: str = "",
@@ -62,8 +62,8 @@ def __init__(
         the :class:`~lightning.app.core.flow.LightningFlow` provided.
 
         Arguments:
-            root: The root LightningFlow component, that defines all the app's nested components, running infinitely.
-                It must define a `run()` method that the app can call.
+            root: The root ``LightningFlow`` or ``LightningWork`` component, that defines all the app's nested
+                 components, running infinitely. It must define a `run()` method that the app can call.
             debug: Whether to activate the Lightning Logger debug mode.
                 This can be helpful when reporting bugs on Lightning repo.
             info: Provide additional info about the app which will be used to update html title,
@@ -89,6 +89,10 @@ def __init__(
         """
 
         self.root_path = root_path  # when running behind a proxy
+
+        if isinstance(root, lightning_app.LightningWork):
+            root = lightning_app.core.flow._RootFlow(root)
+
         _validate_root_flow(root)
         self._root = root
 

diff --git a/src/lightning_app/core/flow.py b/src/lightning_app/core/flow.py
@@ -667,3 +667,15 @@ def configure_api(self):
         under the ``/docs`` route.
         """
         raise NotImplementedError
+
+
+class _RootFlow(LightningFlow):
+    def __init__(self, work):
+        super().__init__()
+        self.work = work
+
+    def run(self):
+        self.work.run()
+
+    def configure_layout(self):
+        return [{"name": "Main", "content": self.work}]
diff --git a/tests/tests_app/core/test_lightning_work.py b/tests/tests_app/core/test_lightning_work.py
@@ -10,6 +10,7 @@
 from lightning_app.runners import MultiProcessRuntime
 from lightning_app.storage import Path
 from lightning_app.testing.helpers import EmptyFlow, EmptyWork, MockQueue
+from lightning_app.testing.testing import LightningTestApp
 from lightning_app.utilities.enum import WorkStageStatus
 from lightning_app.utilities.proxies import ProxyWorkRun, WorkRunner
 
@@ -327,3 +328,20 @@ def run(self, *args, **kwargs):
 
     w = Work()
     w.run()
+
+
+class WorkCounter(LightningWork):
+    def run(self):
+        pass
+
+
+class LightningTestAppWithWork(LightningTestApp):
+    def on_before_run_once(self):
+        if self.root.work.has_succeeded:
+            return True
+        return super().on_before_run_once()
+
+
+def test_lightning_app_with_work():
+    app = LightningTestAppWithWork(WorkCounter())
+    MultiProcessRuntime(app, start_server=False).dispatch()