From ebe784881d6ffcef83a36933bbed6f474951163d Mon Sep 17 00:00:00 2001 From: Yurij Mikhalevich Date: Wed, 14 Dec 2022 20:20:32 +0100 Subject: [PATCH] feature(cli): login flow fixes and improvements (#16052) Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> --- src/lightning_app/cli/lightning_cli.py | 2 +- src/lightning_app/utilities/login.py | 45 +++++++++++++++++--------- 2 files changed, 31 insertions(+), 16 deletions(-) diff --git a/src/lightning_app/cli/lightning_cli.py b/src/lightning_app/cli/lightning_cli.py index f17d4775784d9..19a6c56ff0fc0 100644 --- a/src/lightning_app/cli/lightning_cli.py +++ b/src/lightning_app/cli/lightning_cli.py @@ -205,7 +205,7 @@ def login() -> None: auth.clear() try: - auth._run_server() + auth.authenticate() except ConnectionError: click.echo(f"Unable to connect to {get_lightning_cloud_url()}. Please check your internet connection.") exit(1) diff --git a/src/lightning_app/utilities/login.py b/src/lightning_app/utilities/login.py index 4539ef805eafa..ff04e3b865d41 100644 --- a/src/lightning_app/utilities/login.py +++ b/src/lightning_app/utilities/login.py @@ -4,6 +4,7 @@ import pathlib from dataclasses import dataclass from enum import Enum +from time import sleep from typing import Optional from urllib.parse import urlencode @@ -44,9 +45,11 @@ def __post_init__(self): setattr(self, key.suffix, os.environ.get(key.value, None)) self._with_env_var = bool(self.user_id and self.api_key) # used by authenticate method - if self.api_key and not self.user_id: + if self._with_env_var: + self.save("", self.user_id, self.api_key, self.user_id) + logger.info("Credentials loaded from environment variables") + elif self.api_key or self.user_id: raise ValueError( - f"{Keys.USER_ID.value} is missing from env variables. " "To use env vars for authentication both " f"{Keys.USER_ID.value} and {Keys.API_KEY.value} should be set." ) @@ -135,7 +138,8 @@ def authenticate(self) -> Optional[str]: class AuthServer: - def get_auth_url(self, port: int) -> str: + @staticmethod + def get_auth_url(port: int) -> str: redirect_uri = f"http://localhost:{port}/login-complete" params = urlencode(dict(redirectTo=redirect_uri)) return f"{get_lightning_cloud_url()}/sign-in?{params}" @@ -144,6 +148,7 @@ def login_with_browser(self, auth: Auth) -> None: app = FastAPI() port = find_free_network_port() url = self.get_auth_url(port) + try: # check if server is reachable or catch any network errors requests.head(url) @@ -156,32 +161,42 @@ def login_with_browser(self, auth: Auth) -> None: f"An error occurred with the request. Please report this issue to Lightning Team \n{e}" # E501 ) - logger.info(f"login started for lightning.ai, opening {url}") + logger.info( + "\nAttempting to automatically open the login page in your default browser.\n" + 'If the browser does not open, navigate to the "Keys" tab on your Lightning AI profile page:\n\n' + f"{get_lightning_cloud_url()}/me/keys\n\n" + 'Copy the "Headless CLI Login" command, and execute it in your terminal.\n' + ) click.launch(url) @app.get("/login-complete") async def save_token(request: Request, token="", key="", user_id: str = Query("", alias="userID")): - if token: - auth.save(token=token, username=user_id, user_id=user_id, api_key=key) - logger.info("Authentication Successful") - else: + async def stop_server_once_request_is_done(): + while not await request.is_disconnected(): + sleep(0.25) + server.should_exit = True + + if not token: logger.warn( - "Authentication Failed. This is most likely because you're using an older version of the CLI. \n" # noqa E501 + "Login Failed. This is most likely because you're using an older version of the CLI. \n" # noqa E501 "Please try to update the CLI or open an issue with this information \n" # E501 f"expected token in {request.query_params.items()}" ) + return RedirectResponse( + url=f"{get_lightning_cloud_url()}/cli-login-failed", + background=BackgroundTask(stop_server_once_request_is_done), + ) + + auth.save(token=token, username=user_id, user_id=user_id, api_key=key) + logger.info("Login Successful") # Include the credentials in the redirect so that UI will also be logged in params = urlencode(dict(token=token, key=key, userID=user_id)) return RedirectResponse( - url=f"{get_lightning_cloud_url()}/me/apps?{params}", - # The response background task is being executed right after the server finished writing the response - background=BackgroundTask(stop_server), + url=f"{get_lightning_cloud_url()}/cli-login-successful?{params}", + background=BackgroundTask(stop_server_once_request_is_done), ) - def stop_server(): - server.should_exit = True - server = uvicorn.Server(config=uvicorn.Config(app, port=port, log_level="error")) server.run()