[App] Improve lightning connect experience (#16035)

(cherry picked from commit e522a12)

Author: tchaton

examples/app_installation_commands/app.py

@@ -13,9 +13,14 @@ def run(self):
         print("lmdb successfully installed")
         print("accessing a module in a Work or Flow body works!")
 
-    @property
-    def ready(self) -> bool:
-        return True
+
+class RootFlow(L.LightningFlow):
+    def __init__(self, work):
+        super().__init__()
+        self.work = work
+
+    def run(self):
+        self.work.run()
 
 
 print(f"accessing an object in main code body works!: version={lmdb.version()}")
@@ -24,4 +29,4 @@ def ready(self) -> bool:
 # run on a cloud machine
 compute = L.CloudCompute("cpu")
 worker = YourComponent(cloud_compute=compute)
-app = L.LightningApp(worker)
+app = L.LightningApp(RootFlow(worker))

src/lightning_app/CHANGELOG.md

@@ -11,6 +11,8 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
 
 - Added `Lightning{Flow,Work}.lightningignores` attributes to programmatically ignore files before uploading to the cloud ([#15818](https://github.com/Lightning-AI/lightning/pull/15818))
 
+- Added a progres bar while connecting to an app through the CLI ([#16035](https://github.com/Lightning-AI/lightning/pull/16035))
+
 
 ### Changed
 

src/lightning_app/cli/commands/connection.py

@@ -8,6 +8,7 @@
 import click
 import psutil
 from lightning_utilities.core.imports import package_available
+from rich.progress import Progress
 
 from lightning_app.utilities.cli_helpers import _LightningAppOpenAPIRetriever
 from lightning_app.utilities.cloud import _get_project
@@ -16,15 +17,33 @@
 from lightning_app.utilities.network import LightningClient
 
 _HOME = os.path.expanduser("~")
-_PPID = str(psutil.Process(os.getpid()).ppid())
+_PPID = os.getenv("LIGHTNING_CONNECT_PPID", str(psutil.Process(os.getpid()).ppid()))
 _LIGHTNING_CONNECTION = os.path.join(_HOME, ".lightning", "lightning_connection")
 _LIGHTNING_CONNECTION_FOLDER = os.path.join(_LIGHTNING_CONNECTION, _PPID)
 
 
 @click.argument("app_name_or_id", required=True)
-@click.option("-y", "--yes", required=False, is_flag=True, help="Whether to download the commands automatically.")
-def connect(app_name_or_id: str, yes: bool = False):
-    """Connect to a Lightning App."""
+def connect(app_name_or_id: str):
+    """Connect your local terminal to a running lightning app.
+
+    After connecting, the lightning CLI will respond to commands exposed by the app.
+
+    Example:
+
+    \b
+    # connect to an app named pizza-cooker-123
+    lightning connect pizza-cooker-123
+    \b
+    # this will now show the commands exposed by pizza-cooker-123
+    lightning --help
+    \b
+    # while connected, you can run the cook-pizza command exposed
+    # by pizza-cooker-123.BTW, this should arguably generate an exception :-)
+    lightning cook-pizza --flavor pineapple
+    \b
+    # once done, disconnect and go back to the standard lightning CLI commands
+    lightning disconnect
+    """
     from lightning_app.utilities.commands.base import _download_command
 
     _clean_lightning_connection()
@@ -47,51 +66,64 @@ def connect(app_name_or_id: str, yes: bool = False):
                 click.echo(f"You are already connected to the cloud Lightning App: {app_name_or_id}.")
         else:
             disconnect()
-            connect(app_name_or_id, yes)
+            connect(app_name_or_id)
 
     elif app_name_or_id.startswith("localhost"):
 
-        if app_name_or_id != "localhost":
-            raise Exception("You need to pass localhost to connect to the local Lightning App.")
+        with Progress() as progress_bar:
+            connecting = progress_bar.add_task("[magenta]Setting things up for you...", total=1.0)
 
-        retriever = _LightningAppOpenAPIRetriever(None)
+            if app_name_or_id != "localhost":
+                raise Exception("You need to pass localhost to connect to the local Lightning App.")
 
-        if retriever.api_commands is None:
-            raise Exception(f"The commands weren't found. Is your app {app_name_or_id} running ?")
+            retriever = _LightningAppOpenAPIRetriever(None)
 
-        commands_folder = os.path.join(_LIGHTNING_CONNECTION_FOLDER, "commands")
-        if not os.path.exists(commands_folder):
-            os.makedirs(commands_folder)
+            if retriever.api_commands is None:
+                raise Exception(f"Connection wasn't successful. Is your app {app_name_or_id} running?")
 
-        _write_commands_metadata(retriever.api_commands)
+            increment = 1 / (1 + len(retriever.api_commands))
 
-        with open(os.path.join(commands_folder, "openapi.json"), "w") as f:
-            json.dump(retriever.openapi, f)
+            progress_bar.update(connecting, advance=increment)
 
-        _install_missing_requirements(retriever, yes)
+            commands_folder = os.path.join(_LIGHTNING_CONNECTION_FOLDER, "commands")
+            if not os.path.exists(commands_folder):
+                os.makedirs(commands_folder)
 
-        for command_name, metadata in retriever.api_commands.items():
-            if "cls_path" in metadata:
-                target_file = os.path.join(commands_folder, f"{command_name.replace(' ','_')}.py")
-                _download_command(
-                    command_name,
-                    metadata["cls_path"],
-                    metadata["cls_name"],
-                    None,
-                    target_file=target_file,
-                )
-                repr_command_name = command_name.replace("_", " ")
-                click.echo(f"Storing `{repr_command_name}` at {target_file}")
-            else:
-                with open(os.path.join(commands_folder, f"{command_name}.txt"), "w") as f:
-                    f.write(command_name)
+            _write_commands_metadata(retriever.api_commands)
+
+            with open(os.path.join(commands_folder, "openapi.json"), "w") as f:
+                json.dump(retriever.openapi, f)
 
-        click.echo(f"You can review all the downloaded commands at {commands_folder}")
+            _install_missing_requirements(retriever)
+
+            for command_name, metadata in retriever.api_commands.items():
+                if "cls_path" in metadata:
+                    target_file = os.path.join(commands_folder, f"{command_name.replace(' ','_')}.py")
+                    _download_command(
+                        command_name,
+                        metadata["cls_path"],
+                        metadata["cls_name"],
+                        None,
+                        target_file=target_file,
+                    )
+                else:
+                    with open(os.path.join(commands_folder, f"{command_name}.txt"), "w") as f:
+                        f.write(command_name)
+
+                progress_bar.update(connecting, advance=increment)
 
         with open(connected_file, "w") as f:
             f.write(app_name_or_id + "\n")
 
-        click.echo("You are connected to the local Lightning App.")
+        click.echo("The lightning CLI now responds to app commands. Use 'lightning --help' to see them.")
+        click.echo(" ")
+
+        Popen(
+            f"LIGHTNING_CONNECT_PPID={_PPID} {sys.executable} -m lightning --help",
+            shell=True,
+            stdout=sys.stdout,
+            stderr=sys.stderr,
+        ).wait()
 
     elif matched_connection_path:
 
@@ -101,40 +133,39 @@ def connect(app_name_or_id: str, yes: bool = False):
             commands = os.path.join(_LIGHTNING_CONNECTION_FOLDER, "commands")
             shutil.copytree(matched_commands, commands)
             shutil.copy(matched_connected_file, connected_file)
-            copied_files = [el for el in os.listdir(commands) if os.path.splitext(el)[1] == ".py"]
-            click.echo("Found existing connection, reusing cached commands")
-            for target_file in copied_files:
-                pretty_command_name = os.path.splitext(target_file)[0].replace("_", " ")
-                click.echo(f"Storing `{pretty_command_name}` at {os.path.join(commands, target_file)}")
 
-        click.echo(f"You can review all the commands at {commands}")
+        click.echo("The lightning CLI now responds to app commands. Use 'lightning --help' to see them.")
         click.echo(" ")
-        click.echo(f"You are connected to the cloud Lightning App: {app_name_or_id}.")
 
-    else:
+        Popen(
+            f"LIGHTNING_CONNECT_PPID={_PPID} {sys.executable} -m lightning --help",
+            shell=True,
+            stdout=sys.stdout,
+            stderr=sys.stderr,
+        ).wait()
 
-        retriever = _LightningAppOpenAPIRetriever(app_name_or_id)
+    else:
+        with Progress() as progress_bar:
+            connecting = progress_bar.add_task("[magenta]Setting things up for you...", total=1.0)
+
+            retriever = _LightningAppOpenAPIRetriever(app_name_or_id)
+
+            if not retriever.api_commands:
+                client = LightningClient()
+                project = _get_project(client)
+                apps = client.lightningapp_instance_service_list_lightningapp_instances(project_id=project.project_id)
+                click.echo(
+                    "We didn't find a matching App. Here are the available Apps that you can"
+                    f"connect to {[app.name for app in apps.lightningapps]}."
+                )
+                return
 
-        if not retriever.api_commands:
-            client = LightningClient()
-            project = _get_project(client)
-            apps = client.lightningapp_instance_service_list_lightningapp_instances(project_id=project.project_id)
-            click.echo(
-                "We didn't find a matching App. Here are the available Apps that could be "
-                f"connected to {[app.name for app in apps.lightningapps]}."
-            )
-            return
+            increment = 1 / (1 + len(retriever.api_commands))
 
-        _install_missing_requirements(retriever, yes)
+            progress_bar.update(connecting, advance=increment)
 
-        if not yes:
-            yes = click.confirm(
-                f"The Lightning App `{app_name_or_id}` provides a command-line (CLI). "
-                "Do you want to proceed and install its CLI ?"
-            )
-            click.echo(" ")
+            _install_missing_requirements(retriever)
 
-        if yes:
             commands_folder = os.path.join(_LIGHTNING_CONNECTION_FOLDER, "commands")
             if not os.path.exists(commands_folder):
                 os.makedirs(commands_folder)
@@ -151,26 +182,25 @@ def connect(app_name_or_id: str, yes: bool = False):
                         retriever.app_id,
                         target_file=target_file,
                     )
-                    pretty_command_name = command_name.replace("_", " ")
-                    click.echo(f"Storing `{pretty_command_name}` at {target_file}")
                 else:
                     with open(os.path.join(commands_folder, f"{command_name}.txt"), "w") as f:
                         f.write(command_name)
 
-            click.echo(f"You can review all the downloaded commands at {commands_folder}")
-
-            click.echo(" ")
-            click.echo("The client interface has been successfully installed. ")
-            click.echo("You can now run the following commands:")
-            for command in retriever.api_commands:
-                pretty_command_name = command.replace("_", " ")
-                click.echo(f"    lightning {pretty_command_name}")
+                progress_bar.update(connecting, advance=increment)
 
         with open(connected_file, "w") as f:
             f.write(retriever.app_name + "\n")
             f.write(retriever.app_id + "\n")
+
+        click.echo("The lightning CLI now responds to app commands. Use 'lightning --help' to see them.")
         click.echo(" ")
-        click.echo(f"You are connected to the cloud Lightning App: {app_name_or_id}.")
+
+        Popen(
+            f"LIGHTNING_CONNECT_PPID={_PPID} {sys.executable} -m lightning --help",
+            shell=True,
+            stdout=sys.stdout,
+            stderr=sys.stderr,
+        ).wait()
 
 
 def disconnect(logout: bool = False):
@@ -244,22 +274,37 @@ def _list_app_commands(echo: bool = True) -> List[str]:
         click.echo("The current Lightning App doesn't have commands.")
         return []
 
+    app_info = metadata[command_names[0]].get("app_info", None)
+
+    title, description, on_connect_end = "Lightning", None, None
+    if app_info:
+        title = app_info.get("title")
+        description = app_info.get("description")
+        on_connect_end = app_info.get("on_connect_end")
+
     if echo:
-        click.echo("Usage: lightning [OPTIONS] COMMAND [ARGS]...")
-        click.echo("")
-        click.echo("  --help     Show this message and exit.")
+        click.echo(f"{title} App")
+        if description:
+            click.echo("")
+            click.echo("Description:")
+            if description.endswith("\n"):
+                description = description[:-2]
+            click.echo(f"  {description}")
         click.echo("")
-        click.echo("Lightning App Commands")
+        click.echo("Commands:")
         max_length = max(len(n) for n in command_names)
         for command_name in command_names:
             padding = (max_length + 1 - len(command_name)) * " "
             click.echo(f"  {command_name}{padding}{metadata[command_name].get('description', '')}")
+        if "LIGHTNING_CONNECT_PPID" in os.environ and on_connect_end:
+            if on_connect_end.endswith("\n"):
+                on_connect_end = on_connect_end[:-2]
+            click.echo(on_connect_end)
     return command_names
 
 
 def _install_missing_requirements(
     retriever: _LightningAppOpenAPIRetriever,
-    yes_global: bool = False,
     fail_if_missing: bool = False,
 ):
     requirements = set()
@@ -281,20 +326,15 @@ def _install_missing_requirements(
                 sys.exit(0)
 
             for req in missing_requirements:
-                if not yes_global:
-                    yes = click.confirm(
-                        f"The Lightning App CLI `{retriever.app_id}` requires `{req}`. Do you want to install it ?"
-                    )
-                else:
-                    print(f"Installing missing `{req}` requirement.")
-                    yes = yes_global
-                if yes:
-                    std_out_out = get_logfile("output.log")
-                    with open(std_out_out, "wb") as stdout:
-                        Popen(
-                            f"{sys.executable} -m pip install {req}", shell=True, stdout=stdout, stderr=sys.stderr
-                        ).wait()
-                    print()
+                std_out_out = get_logfile("output.log")
+                with open(std_out_out, "wb") as stdout:
+                    Popen(
+                        f"{sys.executable} -m pip install {req}",
+                        shell=True,
+                        stdout=stdout,
+                        stderr=stdout,
+                    ).wait()
+                os.remove(std_out_out)
 
 
 def _clean_lightning_connection():

src/lightning_app/cli/lightning_cli.py

@@ -76,8 +76,6 @@ def main() -> None:
                 else:
                     message = f"You are connected to the cloud Lightning App: {app_name}."
 
-                click.echo(" ")
-
                 if (len(sys.argv) > 1 and sys.argv[1] in ["-h", "--help"]) or len(sys.argv) == 1:
                     _list_app_commands()
                 else:

src/lightning_app/components/database/server.py

@@ -14,6 +14,7 @@
 from lightning_app.components.database.utilities import _create_database, _Delete, _Insert, _SelectAll, _Update
 from lightning_app.core.work import LightningWork
 from lightning_app.storage import Drive
+from lightning_app.utilities.app_helpers import Logger
 from lightning_app.utilities.imports import _is_sqlmodel_available
 from lightning_app.utilities.packaging.build_config import BuildConfig
 
@@ -23,6 +24,9 @@
     SQLModel = object
 
 
+logger = Logger(__name__)
+
+
 # Required to avoid Uvicorn Server overriding Lightning App signal handlers.
 # Discussions: https://github.com/encode/uvicorn/discussions/1708
 class _DatabaseUvicornServer(uvicorn.Server):
@@ -167,7 +171,7 @@ def store_database(self):
                 drive = Drive("lit://database", component_name=self.name, root_folder=tmpdir)
                 drive.put(os.path.basename(tmp_db_filename))
 
-            print("Stored the database to the Drive.")
+            logger.debug("Stored the database to the Drive.")
         except Exception:
             print(traceback.print_exc())
 

src/lightning_app/runners/multiprocess.py

@@ -82,7 +82,7 @@ def dispatch(self, *args: Any, open_ui: bool = True, **kwargs: Any):
 
                 if is_overridden("configure_commands", self.app.root):
                     commands = _prepare_commands(self.app)
-                    apis += _commands_to_api(commands)
+                    apis += _commands_to_api(commands, info=self.app.info)
 
                 kwargs = dict(
                     apis=apis,