Run Command from App Comments (#15577)

* initial work

* this seems to work well

* added example test

* updated docs & logging

* fixed errors

* fix typing error

* now using the --setup flag to decide if we should execute app comment commands or not

* updated tests

* added tests

* added test to ci

* fixed failing tests

* code review

* updates

Author: rlizzo

.azure/app-cloud-e2e.yml

@@ -67,6 +67,8 @@ jobs:
           name: "collect_failures"
         'App: custom_work_dependencies':
           name: "custom_work_dependencies"
+        'App: installation_commands_app':
+          name: "installation_commands_app"
         'App: drive':
           name: "drive"
         'App: payload':

docs/source-app/levels/basic/hello_components/build_demo.py

@@ -1,6 +1,6 @@
-# pip install streamlit omegaconf scipy
-# pip install torch
 # app.py
+# !pip install streamlit omegaconf scipy
+# !pip install torch
 
 import lightning as L
 import torch

docs/source-app/levels/basic/hello_components/multi_node.py

@@ -1,4 +1,4 @@
-# ! pip install torch
+# !pip install torch
 import lightning as L
 from lightning.app.components import MultiNode
 

docs/source-app/levels/basic/hello_components/pt_multinode.py

@@ -1,4 +1,4 @@
-# ! pip install torch
+# !pip install torch
 import lightning as L
 from lightning.app.components import MultiNode
 import torch

docs/source-app/levels/basic/hello_components/xgboost.py

@@ -1,5 +1,5 @@
-# pip install sklearn xgboost
 # app.py
+# !pip install sklearn xgboost
 
 import lightning as L
 from sklearn import datasets

docs/source-app/levels/basic/hello_components/xgboost_gpu.py

@@ -1,6 +1,6 @@
-# pip install sklearn xgboost
-# conda install py-xgboost-gpu
 # app.py
+# !pip install sklearn xgboost
+# !conda install py-xgboost-gpu
 import lightning as L
 from sklearn import datasets
 from sklearn.model_selection import train_test_split

examples/app_installation_commands/app.py

@@ -0,0 +1,19 @@
+# EXAMPLE COMPONENT: RUN A SCRIPT
+# app.py
+# !echo "I am installing a dependency not declared in a requirements file"
+# !pip install lmdb
+import lmdb
+
+import lightning as L
+
+
+class YourComponent(L.LightningWork):
+    def run(self):
+        print(lmdb.__version__)
+        print("lmdb successfully installed")
+
+
+# run on a cloud machine
+compute = L.CloudCompute("cpu")
+worker = YourComponent(cloud_compute=compute)
+app = L.LightningApp(worker)

src/lightning_app/cli/lightning_cli.py

@@ -239,6 +239,7 @@ def _run_app(
     open_ui: bool,
     env: tuple,
     secret: tuple,
+    run_app_comment_commands: bool,
 ) -> None:
     file = _prepare_file(file)
 
@@ -285,6 +286,7 @@ def on_before_run(*args: Any, **kwargs: Any) -> None:
         env_vars=env_vars,
         secrets=secrets,
         cluster_id=cluster_id,
+        run_app_comment_commands=run_app_comment_commands,
     )
     if runtime_type == RuntimeType.CLOUD:
         click.echo("Application is ready in the cloud")
@@ -322,6 +324,14 @@ def run() -> None:
 @click.option("--env", type=str, default=[], multiple=True, help="Environment variables to be set for the app.")
 @click.option("--secret", type=str, default=[], multiple=True, help="Secret variables to be set for the app.")
 @click.option("--app_args", type=str, default=[], multiple=True, help="Collection of arguments for the app.")
+@click.option(
+    "--setup",
+    "-s",
+    "run_app_comment_commands",
+    is_flag=True,
+    default=False,
+    help="run environment setup commands from the app comments.",
+)
 def run_app(
     file: str,
     cloud: bool,
@@ -334,9 +344,22 @@ def run_app(
     env: tuple,
     secret: tuple,
     app_args: tuple,
+    run_app_comment_commands: bool,
 ) -> None:
     """Run an app from a file."""
-    _run_app(file, cloud, cluster_id, without_server, no_cache, name, blocking, open_ui, env, secret)
+    _run_app(
+        file,
+        cloud,
+        cluster_id,
+        without_server,
+        no_cache,
+        name,
+        blocking,
+        open_ui,
+        env,
+        secret,
+        run_app_comment_commands,
+    )
 
 
 @_main.group(hidden=True)

src/lightning_app/core/constants.py

@@ -59,6 +59,9 @@ def get_lightning_cloud_url() -> str:
     int(os.getenv("ENABLE_MULTIPLE_WORKS_IN_NON_DEFAULT_CONTAINER", "0"))
 )  # This isn't used in the cloud yet.
 
+# env var trigger running setup commands in the app
+ENABLE_APP_COMMENT_COMMAND_EXECUTION = bool(int(os.getenv("ENABLE_APP_COMMENT_COMMAND_EXECUTION", "0")))
+
 
 DEBUG: bool = lightning_cloud.env.DEBUG
 DEBUG_ENABLED = bool(int(os.getenv("LIGHTNING_DEBUG", "0")))

src/lightning_app/runners/__init__.py

@@ -2,6 +2,15 @@
 from lightning_app.runners.multiprocess import MultiProcessRuntime
 from lightning_app.runners.runtime import dispatch, Runtime
 from lightning_app.runners.singleprocess import SingleProcessRuntime
+from lightning_app.utilities.app_commands import run_app_commands
 from lightning_app.utilities.load_app import load_app_from_file
 
-__all__ = ["dispatch", "load_app_from_file", "Runtime", "MultiProcessRuntime", "SingleProcessRuntime", "CloudRuntime"]
+__all__ = [
+    "dispatch",
+    "load_app_from_file",
+    "run_app_commands",
+    "Runtime",
+    "MultiProcessRuntime",
+    "SingleProcessRuntime",
+    "CloudRuntime",
+]