[App] Auto-upgrade / detect environment mis-match from the CLI

docs/source-app/index.rst

@@ -27,7 +27,7 @@ Install Lightning
 
 .. code-block:: bash
 
-   python -m pip install -U lightning
+   pip install -U lightning
 
 
 Or read the :ref:`advanced install <install>` guide.

docs/source-app/installation.rst

@@ -27,12 +27,12 @@ Install with pip
 
     .. code:: bash
 
-        python -m pip install -U lightning
+        pip install -U lightning
 
 .. note::
 
     If you encounter issues during installation use the following to help troubleshoot:
 
-	.. code:: bash
+    .. code:: bash
 
-            pip list | grep lightning
+        pip list | grep lightning

docs/source-app/installation_mac.rst

@@ -19,4 +19,4 @@ Install the ``lightning`` package
         export GRPC_PYTHON_BUILD_SYSTEM_OPENSSL=1
         export GRPC_PYTHON_BUILD_SYSTEM_ZLIB=1
 
-        python -m pip install -U lightning
+        pip install -U lightning

docs/source-app/installation_win.rst

@@ -31,4 +31,4 @@ Install with pip
 
     .. code:: bash
 
-        python -m pip install -U lightning
+        pip install -U lightning

src/lightning_app/__init__.py

@@ -37,6 +37,8 @@
 if module_available("lightning_app.components.demo"):
     from lightning_app.components import demo  # noqa: F401
 
+__package_name__ = "lightning_app".split(".")[0]
+
 _PACKAGE_ROOT = os.path.dirname(__file__)
 _PROJECT_ROOT = os.path.dirname(os.path.dirname(_PACKAGE_ROOT))
 

src/lightning_app/cli/lightning_cli.py

@@ -29,7 +29,12 @@
 from lightning_app.runners.runtime import dispatch
 from lightning_app.runners.runtime_type import RuntimeType
 from lightning_app.utilities.app_helpers import Logger
-from lightning_app.utilities.cli_helpers import _arrow_time_callback, _format_input_env_variables
+from lightning_app.utilities.cli_helpers import (
+    _arrow_time_callback,
+    _check_environment_and_redirect,
+    _check_version_and_upgrade,
+    _format_input_env_variables,
+)
 from lightning_app.utilities.cluster_logs import _cluster_logs_reader
 from lightning_app.utilities.exceptions import _ApiExceptionHandler, LogLinesLimitExceeded
 from lightning_app.utilities.login import Auth
@@ -49,6 +54,12 @@ def get_app_url(runtime_type: RuntimeType, *args: Any, need_credits: bool = Fals
 
 
 def main() -> None:
+    # Enforce running in PATH Python
+    _check_environment_and_redirect()
+
+    # Check for newer versions and upgrade
+    _check_version_and_upgrade()
+
     # 1: Handle connection to a Lightning App.
     if len(sys.argv) > 1 and sys.argv[1] in ("connect", "disconnect", "logout"):
         _main()

src/lightning_app/utilities/cli_helpers.py

@@ -1,16 +1,25 @@
+import functools
 import json
 import os
 import re
+import shutil
+import subprocess
 import sys
 from typing import Dict, Optional
 
 import arrow
 import click
+import packaging
 import requests
 
+from lightning_app import __package_name__, __version__
 from lightning_app.core.constants import APP_SERVER_PORT
+from lightning_app.utilities.app_helpers import Logger
 from lightning_app.utilities.cloud import _get_project
 from lightning_app.utilities.network import LightningClient
+from lightning_app.utilities.packaging.lightning_utils import get_dist_path_if_editable_install
+
+logger = Logger(__name__)
 
 
 def _format_input_env_variables(env_list: tuple) -> Dict[str, str]:
@@ -214,3 +223,97 @@ def _arrow_time_callback(
             return arrow.get(value)
         except (ValueError, TypeError):
             raise click.ClickException(f"cannot parse time {value}")
+
+
+def _is_valid_release(release):
+    version, release = release
+    version = packaging.version.parse(version)
+    if any(r["yanked"] for r in release) or version.is_devrelease or version.is_prerelease:
+        return False
+    return True
+
+
+@functools.lru_cache(maxsize=1)
+def _is_old_version() -> bool:
+    if get_dist_path_if_editable_install(__package_name__):
+        # Always return False in editable mode
+        return False
+    try:
+        response = requests.get(f"https://pypi.org/pypi/{__package_name__}/json")
+        releases = response.json()["releases"]
+        if __version__ not in releases:
+            # Always return False if not installed from PyPI (e.g. dev versions)
+            return False
+        releases = {version: release for version, release in filter(_is_valid_release, releases.items())}
+        sorted_releases = sorted(
+            releases.items(), key=lambda release: release[1][0]["upload_time_iso_8601"], reverse=True
+        )
+        latest_version = sorted_releases[0][0]
+        return False if __version__ == latest_version else latest_version
+    except Exception:
+        # Return False if any exception occurs
+        return False
+
+
+def _redirect_command(executable: str):
+    """Redirect the current lightning CLI call to the given executable."""
+    subprocess.run(
+        [executable, "-m", "lightning"] + sys.argv[1:],
+        env=os.environ,
+    )
+
+    sys.exit()
+
+
+def _check_version_and_upgrade():
+    new_version = _is_old_version()
+    if new_version:
+        prompt = f"A newer version of {__package_name__} is available ({new_version}). Would you like to upgrade?"
+
+        if click.confirm(prompt, default=True):
+            command = f"pip install --upgrade {__package_name__}"
+
+            logger.info(f"⚡ RUN: {command}")
+
+            # Upgrade
+            subprocess.run(
+                [sys.executable, "-m"] + command.split(" "),
+                check=True,
+            )
+
+            # Re-launch
+            _redirect_command(sys.executable)
+    return
+
+
+def _check_environment_and_redirect():
+    env_executable = shutil.which("python")
+
+    if env_executable != sys.executable:
+        logger.info(
+            "Lightning is running from outside your current environment. " "Switching to your current environment."
+        )
+
+        process = subprocess.run(
+            [env_executable, "-m", "lightning", "--version"],
+            capture_output=True,
+            text=True,
+        )
+
+        if "No module named lightning" in process.stderr:
+            prompt = f"The {__package_name__} package is not installed. " f"Would you like to install it? [Y/n (exit)]"
+
+            if click.confirm(prompt, default=True, show_default=False):
+                command = f"pip install {__package_name__}"
+
+                logger.info(f"⚡ RUN: {command}")
+
+                subprocess.run(
+                    [env_executable, "-m"] + command.split(" "),
+                    check=True,
+                )
+            else:
+                sys.exit()
+
+        _redirect_command(env_executable)
+    return