Fix: App comment command execution sequencing (#15615)

* fixed condition in which app command comments would not execute before the flow is dispatched

* remove debug print statment

* updated example code

* fix test

* updates to test

Co-authored-by: William Falcon <waf2107@columbia.edu>
Co-authored-by: Jirka Borovec <6035284+Borda@users.noreply.github.com>

examples/app_installation_commands/app.py

@@ -9,8 +9,12 @@ import lightning as L
 
 class YourComponent(L.LightningWork):
     def run(self):
-        print(lmdb.__version__)
+        print(lmdb.version())
         print("lmdb successfully installed")
+        print("accessing a module in a Work or Flow body works!")
+
+
+print(f"accessing an object in main code body works!: version={lmdb.version()}")
 
 
 # run on a cloud machine

src/lightning_app/cli/lightning_cli.py

@@ -30,9 +30,10 @@ from lightning_app.cli.lightning_cli_create import create
 from lightning_app.cli.lightning_cli_delete import delete
 from lightning_app.cli.lightning_cli_list import get_list
 from lightning_app.cli.lightning_cli_remove import cli_remove
-from lightning_app.core.constants import DEBUG, get_lightning_cloud_url
+from lightning_app.core.constants import DEBUG, ENABLE_APP_COMMENT_COMMAND_EXECUTION, get_lightning_cloud_url
 from lightning_app.runners.runtime import dispatch
 from lightning_app.runners.runtime_type import RuntimeType
+from lightning_app.utilities.app_commands import run_app_commands
 from lightning_app.utilities.app_helpers import Logger
 from lightning_app.utilities.cli_helpers import (
     _arrow_time_callback,
@@ -261,6 +262,9 @@ def _run_app(
                 "Secrets can only be used for apps running in cloud. "
                 "Using the option --secret in local execution is not supported."
             )
+        if ENABLE_APP_COMMENT_COMMAND_EXECUTION or run_app_comment_commands:
+            if file is not None:
+                run_app_commands(str(file))
 
     env_vars = _format_input_env_variables(env)
     os.environ.update(env_vars)

src/lightning_app/runners/multiprocess.py

@@ -5,11 +5,9 @@ from typing import Any, Callable, Optional, Union
 
 from lightning_app.api.http_methods import _add_tags_to_api, _validate_api
 from lightning_app.core.api import start_server
-from lightning_app.core.constants import ENABLE_APP_COMMENT_COMMAND_EXECUTION
 from lightning_app.runners.backends import Backend
 from lightning_app.runners.runtime import Runtime
 from lightning_app.storage.orchestrator import StorageOrchestrator
-from lightning_app.utilities.app_commands import run_app_commands
 from lightning_app.utilities.app_helpers import is_overridden
 from lightning_app.utilities.commands.base import _commands_to_api, _prepare_commands
 from lightning_app.utilities.component import _set_flow_context, _set_frontend_context
@@ -30,10 +28,6 @@ class MultiProcessRuntime(Runtime):
 
     def dispatch(self, *args: Any, on_before_run: Optional[Callable] = None, **kwargs: Any):
         """Method to dispatch and run the LightningApp."""
-        if ENABLE_APP_COMMENT_COMMAND_EXECUTION or self.run_app_comment_commands:
-            if self.entrypoint_file is not None:
-                run_app_commands(str(self.entrypoint_file))
-
         try:
             _set_flow_context()
             self.app.backend = self.backend

src/lightning_app/runners/singleprocess.py

@@ -2,10 +2,8 @@ import multiprocessing as mp
 from typing import Any, Callable, Optional
 
 from lightning_app.core.api import start_server
-from lightning_app.core.constants import ENABLE_APP_COMMENT_COMMAND_EXECUTION
 from lightning_app.core.queues import QueuingSystem
 from lightning_app.runners.runtime import Runtime
-from lightning_app.utilities.app_commands import run_app_commands
 from lightning_app.utilities.load_app import extract_metadata_from_app
 
 
@@ -17,11 +15,6 @@ class SingleProcessRuntime(Runtime):
 
     def dispatch(self, *args, on_before_run: Optional[Callable] = None, **kwargs: Any):
         """Method to dispatch and run the LightningApp."""
-
-        if ENABLE_APP_COMMENT_COMMAND_EXECUTION or self.run_app_comment_commands:
-            if self.entrypoint_file is not None:
-                run_app_commands(str(self.entrypoint_file))
-
         queue = QueuingSystem.SINGLEPROCESS
 
         self.app.delta_queue = queue.get_delta_queue()

tests/tests_app_examples/test_installation_commands_app.py

@@ -1,5 +1,4 @@
 import os
-from datetime import time
 
 import pytest
 from tests_app import _PROJECT_ROOT
@@ -7,9 +6,9 @@ from tests_app import _PROJECT_ROOT
 from lightning_app.testing.testing import run_app_in_cloud
 
 
-@pytest.mark.skip(reason="temporarily disabled until backend release")
 @pytest.mark.cloud
 def test_installation_commands_app_example_cloud() -> None:
+    # This is expected to pass, since the "setup" flag is passed
     with run_app_in_cloud(
         os.path.join(_PROJECT_ROOT, "examples/app_installation_commands"),
         app_name="app.py",
@@ -21,4 +20,3 @@ def test_installation_commands_app_example_cloud() -> None:
             for log in fetch_logs(["work"]):
                 if "lmdb successfully installed" in log:
                     has_logs = True
-            time.sleep(1)