[App] Improve lightning connect experience

src/lightning_app/cli/commands/connection.py

@@ -3,11 +3,14 @@
 import shutil
 import sys
 from subprocess import Popen
+from threading import Event, Thread
+from time import sleep, time
 from typing import List, Optional, Tuple
 
 import click
 import psutil
 from lightning_utilities.core.imports import package_available
+from rich.progress import Progress, Task
 
 from lightning_app.utilities.cli_helpers import _LightningAppOpenAPIRetriever
 from lightning_app.utilities.cloud import _get_project
@@ -16,14 +19,14 @@
 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)
+_PROGRESS_TOTAL = 60
 
 
 @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):
+def connect(app_name_or_id: str):
     """Connect to a Lightning App."""
     from lightning_app.utilities.commands.base import _download_command
 
@@ -47,51 +50,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:
+            event = Event()
+            connecting = progress_bar.add_task("[magenta]Setting things up for you...", total=_PROGRESS_TOTAL)
+            thread = Thread(target=update_progresss, args=[event, progress_bar, connecting], daemon=True)
+            thread.start()
 
-        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"The commands weren't found. Is your app {app_name_or_id} running ?")
+
+            commands_folder = os.path.join(_LIGHTNING_CONNECTION_FOLDER, "commands")
+            if not os.path.exists(commands_folder):
+                os.makedirs(commands_folder)
 
-        _write_commands_metadata(retriever.api_commands)
+            _write_commands_metadata(retriever.api_commands)
 
-        with open(os.path.join(commands_folder, "openapi.json"), "w") as f:
-            json.dump(retriever.openapi, f)
+            with open(os.path.join(commands_folder, "openapi.json"), "w") as f:
+                json.dump(retriever.openapi, f)
 
-        _install_missing_requirements(retriever, yes)
+            _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,
-                )
-                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)
+            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)
 
-        click.echo(f"You can review all the downloaded commands at {commands_folder}")
+            event.set()
+            thread.join()
 
         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 respond 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 +117,38 @@ 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 respond 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:
 
-        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 could be "
-                f"connected to {[app.name for app in apps.lightningapps]}."
-            )
-            return
+        Popen(
+            f"LIGHTNING_CONNECT_PPID={_PPID} {sys.executable} -m lightning --help",
+            shell=True,
+            stdout=sys.stdout,
+            stderr=sys.stderr,
+        ).wait()
 
-        _install_missing_requirements(retriever, yes)
+    else:
+        with Progress() as progress_bar:
+            event = Event()
+            connecting = progress_bar.add_task("[magenta]Setting things up for you...", total=_PROGRESS_TOTAL)
+            thread = Thread(target=update_progresss, args=[event, progress_bar, connecting], daemon=True)
+            thread.start()
+
+            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 could be "
+                    f"connected to {[app.name for app in apps.lightningapps]}."
+                )
+                return
 
-        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 +165,26 @@ 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}")
+            event.set()
+            thread.join()
 
         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 respond 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 +258,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_after_connect = "Lightning", None, None
+    if app_info:
+        title = app_info["title"]
+        description = app_info["description"]
+        on_after_connect = app_info["on_after_connect"]
+
     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_after_connect:
+            if on_after_connect.endswith("\n"):
+                on_after_connect = on_after_connect[:-2]
+            click.echo(on_after_connect)
     return command_names
 
 
 def _install_missing_requirements(
     retriever: _LightningAppOpenAPIRetriever,
-    yes_global: bool = False,
     fail_if_missing: bool = False,
 ):
     requirements = set()
@@ -281,20 +310,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():
@@ -332,3 +356,17 @@ def _scan_lightning_connections(app_name_or_id):
             return connection_path
 
     return None
+
+
+def update_progresss(exit_event: Event, progress_bar: Progress, task: Task) -> None:
+    t0 = time()
+    while not exit_event.is_set():
+        sleep(0.5)
+        progress_bar.update(task, advance=0.5)
+
+    # Note: This is required to make progress feel more naturally progressing.
+    remaning = _PROGRESS_TOTAL - (time() - t0)
+    num_updates = 10
+    for _ in range(num_updates):
+        progress_bar.update(task, advance=remaning / float(num_updates))
+        sleep(0.2)

src/lightning_app/cli/lightning_cli.py

@@ -77,8 +77,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,

src/lightning_app/utilities/cli_helpers.py

@@ -69,6 +69,7 @@ def _get_metadata_from_openapi(paths: Dict, path: str):
     cls_name = paths[path]["post"].get("cls_name", None)
     description = paths[path]["post"].get("description", None)
     requirements = paths[path]["post"].get("requirements", None)
+    app_info = paths[path]["post"].get("app_info", None)
 
     metadata = {"tag": tag, "parameters": {}}
 
@@ -84,6 +85,9 @@ def _get_metadata_from_openapi(paths: Dict, path: str):
     if description:
         metadata["requirements"] = requirements
 
+    if app_info:
+        metadata["app_info"] = app_info
+
     if not parameters:
         return metadata