[App] Auto-upgrade / detect environment mis-match from the CLI (#15434)

* Add auto-upgrade from the CLI and check for current env

* No longer require `python -m` in docs

* Tabs -> spaces

* Ignore pre-releases

* Test + docs

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,17 +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
 
+logger = Logger(__name__)
+
 
 def _format_input_env_variables(env_list: tuple) -> Dict[str, str]:
     """
@@ -214,3 +222,108 @@ 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 _get_newer_version() -> Optional[str]:
+    """Check PyPI for newer versions of ``lightning``, returning the newest version if different from the current
+    or ``None`` otherwise."""
+    if packaging.version.parse(__version__).is_prerelease:
+        return None
+    try:
+        response = requests.get(f"https://pypi.org/pypi/{__package_name__}/json")
+        releases = response.json()["releases"]
+        if __version__ not in releases:
+            # Always return None if not installed from PyPI (e.g. dev versions)
+            return None
+        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 None if __version__ == latest_version else latest_version
+    except Exception:
+        # Return None if any exception occurs
+        return "err"
+
+
+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():
+    """Checks that the current version of ``lightning`` is the latest on PyPI.
+
+    If not, prompt the user to upgrade ``lightning`` for them and re-run the current call in the new version.
+    """
+    new_version = _get_newer_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():
+    """Checks that the current ``sys.executable`` is the same as the executable resolved from the current
+    environment.
+
+    If not, this utility tries to redirect the ``lightning`` call to the environment executable (prompting the user to
+    install lightning for them there if needed).
+    """
+    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. 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

tests/tests_app/utilities/test_cli_helpers.py

@@ -1,9 +1,10 @@
-from unittest.mock import Mock
+from unittest.mock import Mock, patch
 
 import arrow
 import pytest
 
-from lightning_app.utilities.cli_helpers import _arrow_time_callback, _format_input_env_variables
+import lightning_app
+from lightning_app.utilities.cli_helpers import _arrow_time_callback, _format_input_env_variables, _get_newer_version
 
 
 def test_format_input_env_variables():
@@ -62,3 +63,51 @@ def test_arrow_time_callback():
 
     with pytest.raises(Exception, match="cannot parse time 1 time unit ago"):
         _arrow_time_callback(Mock(), Mock(), "1 time unit ago")
+
+
+@pytest.mark.parametrize(
+    "releases, current_version, newer_version",
+    [
+        (
+            {
+                "1.0.0": [{"upload_time_iso_8601": "2022-09-10", "yanked": False}],
+                "2.0.0": [{"upload_time_iso_8601": "2022-11-01", "yanked": False}],
+            },
+            "1.0.0",
+            "2.0.0",
+        ),
+        (
+            {
+                "1.0.0": [{"upload_time_iso_8601": "2022-09-10", "yanked": False}],
+                "2.0.0": [{"upload_time_iso_8601": "2022-11-01", "yanked": True}],
+            },
+            "1.0.0",
+            None,
+        ),
+        (
+            {
+                "1.0.0": [{"upload_time_iso_8601": "2022-09-10", "yanked": False}],
+                "2.0.0rc0": [{"upload_time_iso_8601": "2022-11-01", "yanked": False}],
+            },
+            "1.0.0",
+            None,
+        ),
+        (
+            {
+                "2.0.0": [{"upload_time_iso_8601": "2022-11-01", "yanked": False}],
+            },
+            "1.0.0dev",
+            None,
+        ),
+        ({"1.0.0": "this wil trigger an error"}, "1.0.0", "err"),
+        ({}, "1.0.0rc0", None),
+    ],
+)
+@patch("lightning_app.utilities.cli_helpers.requests")
+def test_get_newer_version(mock_requests, releases, current_version, newer_version):
+    mock_requests.get().json.return_value = {"releases": releases}
+
+    lightning_app.utilities.cli_helpers.__version__ = current_version
+
+    _get_newer_version.cache_clear()
+    assert _get_newer_version() == newer_version