From 0299d4e502c0c116abc7ea9c9af640d2086de001 Mon Sep 17 00:00:00 2001 From: Ethan Harris Date: Wed, 2 Nov 2022 21:13:36 +0000 Subject: [PATCH] [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 (cherry picked from commit 555257a4ba0e8b7b7f76687b0c6209962de1602f) --- docs/source-app/index.rst | 2 +- docs/source-app/installation.rst | 6 +- docs/source-app/installation_mac.rst | 2 +- docs/source-app/installation_win.rst | 2 +- src/lightning_app/__init__.py | 2 + src/lightning_app/cli/lightning_cli.py | 13 +- src/lightning_app/utilities/cli_helpers.py | 113 ++++++++++++++++++ tests/tests_app/utilities/test_cli_helpers.py | 53 +++++++- 8 files changed, 184 insertions(+), 9 deletions(-) diff --git a/docs/source-app/index.rst b/docs/source-app/index.rst index c818aec936bcf..3e635ddac4b6d 100644 --- a/docs/source-app/index.rst +++ b/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 ` guide. diff --git a/docs/source-app/installation.rst b/docs/source-app/installation.rst index 0faa50216e9af..61a8cf9ce73ca 100644 --- a/docs/source-app/installation.rst +++ b/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 diff --git a/docs/source-app/installation_mac.rst b/docs/source-app/installation_mac.rst index d6274be24fdf5..f5dd13bd0a444 100644 --- a/docs/source-app/installation_mac.rst +++ b/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 diff --git a/docs/source-app/installation_win.rst b/docs/source-app/installation_win.rst index ff08cc1945da3..a3c548840e21a 100644 --- a/docs/source-app/installation_win.rst +++ b/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 diff --git a/src/lightning_app/__init__.py b/src/lightning_app/__init__.py index e64e2269aa39e..893ed4da9f62c 100644 --- a/src/lightning_app/__init__.py +++ b/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)) diff --git a/src/lightning_app/cli/lightning_cli.py b/src/lightning_app/cli/lightning_cli.py index 5d814d4902f28..b4509ca58f20f 100644 --- a/src/lightning_app/cli/lightning_cli.py +++ b/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() diff --git a/src/lightning_app/utilities/cli_helpers.py b/src/lightning_app/utilities/cli_helpers.py index 492d87f8caeb3..4b404ebf56661 100644 --- a/src/lightning_app/utilities/cli_helpers.py +++ b/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 diff --git a/tests/tests_app/utilities/test_cli_helpers.py b/tests/tests_app/utilities/test_cli_helpers.py index 4711ffeddbfde..4ebb3ddc4f0ae 100644 --- a/tests/tests_app/utilities/test_cli_helpers.py +++ b/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