From 88e089ea4ec9dae2eeb87779eda1c3fd0a533633 Mon Sep 17 00:00:00 2001 From: thomas chaton Date: Wed, 8 Feb 2023 18:41:05 +0000 Subject: [PATCH] [App] Enable to register data connections (#16670) Co-authored-by: thomas --- pyproject.toml | 3 +- requirements/app/base.txt | 2 +- src/lightning/app/CHANGELOG.md | 6 ++ .../app/cli/commands/app_commands.py | 2 +- src/lightning/app/cli/commands/cd.py | 2 +- src/lightning/app/cli/commands/ls.py | 5 +- src/lightning/app/cli/connect/__init__.py | 0 .../connection.py => connect/app.py} | 8 +- src/lightning/app/cli/connect/data.py | 92 +++++++++++++++++++ src/lightning/app/cli/lightning_cli.py | 32 +++++-- .../public/test_commands_and_api.py | 4 +- tests/tests_app/cli/test_connect.py | 42 ++++----- tests/tests_app/cli/test_connect_data.py | 60 ++++++++++++ tests/tests_app/utilities/test_commands.py | 6 +- 14 files changed, 218 insertions(+), 46 deletions(-) create mode 100644 src/lightning/app/cli/connect/__init__.py rename src/lightning/app/cli/{commands/connection.py => connect/app.py} (98%) create mode 100644 src/lightning/app/cli/connect/data.py create mode 100644 tests/tests_app/cli/test_connect_data.py diff --git a/pyproject.toml b/pyproject.toml index 7ddaebd84d..55b60b8fd2 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -105,13 +105,14 @@ module = [ "lightning.app.api.http_methods", "lightning.app.api.request_types", "lightning.app.cli.commands.app_commands", - "lightning.app.cli.commands.connection", "lightning.app.cli.commands.lightning_cli", "lightning.app.cli.commands.cmd_install", + "lightning.app.cli.connect.app", "lightning.app.cli.commands.pwd", "lightning.app.cli.commands.ls", "lightning.app.cli.commands.cp", "lightning.app.cli.commands.cd", + "lightning.app.cli.connect.data", "lightning.app.cli.cmd_install", "lightning.app.components.database.client", "lightning.app.components.database.server", diff --git a/requirements/app/base.txt b/requirements/app/base.txt index 8647f94d2a..87b1c1f0b5 100644 --- a/requirements/app/base.txt +++ b/requirements/app/base.txt @@ -1,4 +1,4 @@ -lightning-cloud>=0.5.22 +lightning-cloud>=0.5.24 packaging typing-extensions>=4.0.0, <=4.4.0 deepdiff>=5.7.0, <6.2.4 diff --git a/src/lightning/app/CHANGELOG.md b/src/lightning/app/CHANGELOG.md index f1a44c074c..13b2c6db51 100644 --- a/src/lightning/app/CHANGELOG.md +++ b/src/lightning/app/CHANGELOG.md @@ -22,6 +22,9 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/). * `pwd`: Return the current folder in your Cloud Platform Filesystem * `cp`: Copy files between your Cloud Platform Filesystem and local filesystem +- Added `lightning connect data` to register data connection to s3 buckets ([#16670](https://github.com/Lightning-AI/lightning/pull/16670)) + + ### Changed - Changed the default `LightningClient(retry=False)` to `retry=True` ([#16382](https://github.com/Lightning-AI/lightning/pull/16382)) @@ -33,6 +36,9 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/). - Renamed `lightning.app.components.LiteMultiNode` to `lightning.app.components.FabricMultiNode` ([#16505](https://github.com/Lightning-AI/lightning/pull/16505)) +- Changed the command `lightning connect` to `lightning connect app` for consistency ([#16670](https://github.com/Lightning-AI/lightning/pull/16670)) + + ### Deprecated - diff --git a/src/lightning/app/cli/commands/app_commands.py b/src/lightning/app/cli/commands/app_commands.py index 418ae29d5c..2251ef74f9 100644 --- a/src/lightning/app/cli/commands/app_commands.py +++ b/src/lightning/app/cli/commands/app_commands.py @@ -18,7 +18,7 @@ from typing import Dict, Optional import requests -from lightning.app.cli.commands.connection import ( +from lightning.app.cli.connect.app import ( _clean_lightning_connection, _install_missing_requirements, _resolve_command_path, diff --git a/src/lightning/app/cli/commands/cd.py b/src/lightning/app/cli/commands/cd.py index f25a760702..80f1974755 100644 --- a/src/lightning/app/cli/commands/cd.py +++ b/src/lightning/app/cli/commands/cd.py @@ -21,7 +21,7 @@ from rich.spinner import Spinner from rich.text import Text from lightning.app.cli.commands import ls -from lightning.app.cli.commands.connection import _LIGHTNING_CONNECTION_FOLDER +from lightning.app.cli.connect.app import _LIGHTNING_CONNECTION_FOLDER from lightning.app.utilities.app_helpers import Logger from lightning.app.utilities.cli_helpers import _error_and_exit diff --git a/src/lightning/app/cli/commands/ls.py b/src/lightning/app/cli/commands/ls.py index 1339ec764d..4b3c2a8414 100644 --- a/src/lightning/app/cli/commands/ls.py +++ b/src/lightning/app/cli/commands/ls.py @@ -26,7 +26,7 @@ from rich.live import Live from rich.spinner import Spinner from rich.text import Text -from lightning.app.cli.commands.connection import _LIGHTNING_CONNECTION_FOLDER +from lightning.app.cli.connect.app import _LIGHTNING_CONNECTION_FOLDER from lightning.app.utilities.app_helpers import Logger from lightning.app.utilities.cli_helpers import _error_and_exit from lightning.app.utilities.network import LightningClient @@ -44,8 +44,7 @@ def ls(path: Optional[str] = None, print: bool = True, use_live: bool = True) -> from lightning.app.cli.commands.cd import _CD_FILE if sys.platform == "win32": - print("`ls` isn't supported on windows. Open an issue on Github.") - sys.exit(0) + _error_and_exit("`ls` isn't supported on windows. Open an issue on Github.") root = "/" diff --git a/src/lightning/app/cli/connect/__init__.py b/src/lightning/app/cli/connect/__init__.py new file mode 100644 index 0000000000..e69de29bb2 diff --git a/src/lightning/app/cli/commands/connection.py b/src/lightning/app/cli/connect/app.py similarity index 98% rename from src/lightning/app/cli/commands/connection.py rename to src/lightning/app/cli/connect/app.py index b15fa6d03d..4f181da033 100644 --- a/src/lightning/app/cli/commands/connection.py +++ b/src/lightning/app/cli/connect/app.py @@ -37,7 +37,7 @@ _LIGHTNING_CONNECTION_FOLDER = os.path.join(_LIGHTNING_CONNECTION, _PPID) @click.argument("app_name_or_id", required=True) -def connect(app_name_or_id: str): +def connect_app(app_name_or_id: str): """Connect your local terminal to a running lightning app. After connecting, the lightning CLI will respond to commands exposed by the app. @@ -79,8 +79,8 @@ def connect(app_name_or_id: str): else: click.echo(f"You are already connected to the cloud Lightning App: {app_name_or_id}.") else: - disconnect() - connect(app_name_or_id) + disconnect_app() + connect_app(app_name_or_id) elif app_name_or_id.startswith("localhost"): @@ -217,7 +217,7 @@ def connect(app_name_or_id: str): ).wait() -def disconnect(logout: bool = False): +def disconnect_app(logout: bool = False): """Disconnect from an App.""" _clean_lightning_connection() diff --git a/src/lightning/app/cli/connect/data.py b/src/lightning/app/cli/connect/data.py new file mode 100644 index 0000000000..1c9efcd1f5 --- /dev/null +++ b/src/lightning/app/cli/connect/data.py @@ -0,0 +1,92 @@ +# Copyright The Lightning team. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import sys + +import click +import rich +from lightning_cloud.openapi import ProjectIdDataConnectionsBody +from rich.live import Live +from rich.spinner import Spinner +from rich.text import Text + +from lightning.app.utilities.app_helpers import Logger +from lightning.app.utilities.cli_helpers import _error_and_exit +from lightning.app.utilities.cloud import _get_project +from lightning.app.utilities.network import LightningClient + +logger = Logger(__name__) + + +@click.argument("name", required=True) +@click.argument("region", required=True) +@click.argument("source", required=True) +@click.argument("destination", required=False) +@click.argument("project_name", required=False) +def connect_data( + name: str, + region: str, + source: str, + destination: str = "", + project_name: str = "", +) -> None: + """Create a new data connection.""" + + if sys.platform == "win32": + _error_and_exit("Data connection isn't supported on windows. Open an issue on Github.") + + with Live(Spinner("point", text=Text("pending...", style="white")), transient=True) as live: + + live.stop() + + client = LightningClient() + projects = client.projects_service_list_memberships() + + project_id = None + + for project in projects.memberships: + + if project.name == project_name: + project_id = project.project_id + break + + if project_id is None: + project_id = _get_project(client).project_id + + if not source.startswith("s3://"): + return _error_and_exit( + "Only public S3 folders are supported for now. Please, open a Github issue with your use case." + ) + + try: + _ = client.data_connection_service_create_data_connection( + body=ProjectIdDataConnectionsBody( + name=name, + region=region, + source=source, + destination=destination, + ), + project_id=project_id, + ) + + # Note: Expose through lightning show data {DATA_NAME} + # response = client.data_connection_service_list_data_connection_artifacts( + # project_id=project_id, + # id=response.id, + # ) + # print(response) + except Exception: + _error_and_exit("The data connection creation failed.") + + rich.print(f"[green]Succeeded[/green]: You have created a new data connection {name}.") diff --git a/src/lightning/app/cli/lightning_cli.py b/src/lightning/app/cli/lightning_cli.py index 56f8ced08b..0a7b0484e7 100644 --- a/src/lightning/app/cli/lightning_cli.py +++ b/src/lightning/app/cli/lightning_cli.py @@ -34,16 +34,17 @@ from lightning.app.cli.cmd_apps import _AppManager from lightning.app.cli.cmd_clusters import AWSClusterManager from lightning.app.cli.commands.app_commands import _run_app_command from lightning.app.cli.commands.cd import cd -from lightning.app.cli.commands.connection import ( - _list_app_commands, - _retrieve_connection_to_an_app, - connect, - disconnect, -) from lightning.app.cli.commands.cp import cp from lightning.app.cli.commands.logs import logs from lightning.app.cli.commands.ls import ls from lightning.app.cli.commands.pwd import pwd +from lightning.app.cli.connect.app import ( + _list_app_commands, + _retrieve_connection_to_an_app, + connect_app, + disconnect_app, +) +from lightning.app.cli.connect.data import connect_data from lightning.app.cli.lightning_cli_create import create from lightning.app.cli.lightning_cli_delete import delete from lightning.app.cli.lightning_cli_list import get_list @@ -121,8 +122,21 @@ def show() -> None: pass -_main.command()(connect) -_main.command()(disconnect) +@_main.group() +def connect() -> None: + """Connect apps and data.""" + pass + + +@_main.group() +def disconnect() -> None: + """Disconnect apps.""" + pass + + +connect.command("app")(connect_app) +disconnect.command("app")(disconnect_app) +connect.command("data")(connect_data) _main.command()(ls) _main.command()(cd) _main.command()(cp) @@ -240,7 +254,7 @@ def login() -> None: def logout() -> None: """Log out of your lightning.ai account.""" Auth().clear() - disconnect(logout=True) + disconnect_app(logout=True) def _run_app( diff --git a/tests/integrations_app/public/test_commands_and_api.py b/tests/integrations_app/public/test_commands_and_api.py index ddafbca01d..16e9b071b8 100644 --- a/tests/integrations_app/public/test_commands_and_api.py +++ b/tests/integrations_app/public/test_commands_and_api.py @@ -19,11 +19,11 @@ def test_commands_and_api_example_cloud() -> None: ): # Connect to the App and send the first & second command with the client # Requires to be run within the same process. - cmd_1 = f"python -m lightning connect {app_name}" + cmd_1 = f"python -m lightning connect app {app_name}" cmd_2 = "python -m lightning command with client --name=this" cmd_3 = "python -m lightning command without client --name=is" cmd_4 = "python -m lightning command without client --name=awesome" - cmd_5 = "lightning disconnect" + cmd_5 = "lightning disconnect app" process = Popen(" && ".join([cmd_1, cmd_2, cmd_3, cmd_4, cmd_5]), shell=True) process.wait() diff --git a/tests/tests_app/cli/test_connect.py b/tests/tests_app/cli/test_connect.py index 9837f907b3..524a485b2d 100644 --- a/tests/tests_app/cli/test_connect.py +++ b/tests/tests_app/cli/test_connect.py @@ -7,12 +7,12 @@ import psutil import pytest from lightning.app import _PROJECT_ROOT -from lightning.app.cli.commands.connection import ( +from lightning.app.cli.connect.app import ( _list_app_commands, _resolve_command_path, _retrieve_connection_to_an_app, - connect, - disconnect, + connect_app, + disconnect_app, ) from lightning.app.utilities import cli_helpers from lightning.app.utilities.commands import base @@ -20,18 +20,18 @@ from lightning.app.utilities.commands import base def monkeypatch_connection(monkeypatch, tmpdir, ppid): connection_path = os.path.join(tmpdir, ppid) - monkeypatch.setattr("lightning.app.cli.commands.connection._clean_lightning_connection", MagicMock()) - monkeypatch.setattr("lightning.app.cli.commands.connection._PPID", ppid) - monkeypatch.setattr("lightning.app.cli.commands.connection._LIGHTNING_CONNECTION", tmpdir) - monkeypatch.setattr("lightning.app.cli.commands.connection._LIGHTNING_CONNECTION_FOLDER", connection_path) + monkeypatch.setattr("lightning.app.cli.connect.app._clean_lightning_connection", MagicMock()) + monkeypatch.setattr("lightning.app.cli.connect.app._PPID", ppid) + monkeypatch.setattr("lightning.app.cli.connect.app._LIGHTNING_CONNECTION", tmpdir) + monkeypatch.setattr("lightning.app.cli.connect.app._LIGHTNING_CONNECTION_FOLDER", connection_path) return connection_path def test_connect_disconnect_local(tmpdir, monkeypatch): - disconnect() + disconnect_app() with pytest.raises(Exception, match="Connection wasn't successful. Is your app localhost running ?"): - connect("localhost") + connect_app("localhost") with open(os.path.join(os.path.dirname(__file__), "jsons/connect_1.json")) as f: data = json.load(f) @@ -43,7 +43,7 @@ def test_connect_disconnect_local(tmpdir, monkeypatch): messages = [] - disconnect() + disconnect_app() def fn(msg): messages.append(msg) @@ -54,21 +54,21 @@ def test_connect_disconnect_local(tmpdir, monkeypatch): response.status_code = 200 response.json.return_value = data monkeypatch.setattr(cli_helpers.requests, "get", MagicMock(return_value=response)) - connect("localhost") + connect_app("localhost") assert _retrieve_connection_to_an_app() == ("localhost", None) command_path = _resolve_command_path("nested_command") assert not os.path.exists(command_path) command_path = _resolve_command_path("command_with_client") assert os.path.exists(command_path) messages = [] - connect("localhost") + connect_app("localhost") assert messages == ["You are connected to the local Lightning App."] messages = [] - disconnect() + disconnect_app() assert messages == ["You are disconnected from the local Lightning App."] messages = [] - disconnect() + disconnect_app() assert messages == [ "You aren't connected to any Lightning App. Please use `lightning connect app_name_or_id` to connect to one." ] @@ -77,7 +77,7 @@ def test_connect_disconnect_local(tmpdir, monkeypatch): def test_connect_disconnect_cloud(tmpdir, monkeypatch): - disconnect() + disconnect_app() ppid_1 = str(psutil.Process(os.getpid()).ppid()) ppid_2 = "222" @@ -132,7 +132,7 @@ def test_connect_disconnect_cloud(tmpdir, monkeypatch): with open(data["paths"]["/command/command_with_client"]["post"]["cls_path"], "rb") as f: response.content = f.read() - connect("example") + connect_app("example") assert _retrieve_connection_to_an_app() == ("example", "1234") commands = _list_app_commands() assert commands == ["command with client", "command without client", "nested command"] @@ -141,27 +141,27 @@ def test_connect_disconnect_cloud(tmpdir, monkeypatch): command_path = _resolve_command_path("command_with_client") assert os.path.exists(command_path) messages = [] - connect("example") + connect_app("example") assert messages == ["You are already connected to the cloud Lightning App: example."] _ = monkeypatch_connection(monkeypatch, tmpdir, ppid=ppid_2) messages = [] - connect("example") + connect_app("example") assert "The lightning CLI now responds to app commands" in messages[0] messages = [] - disconnect() + disconnect_app() assert messages == ["You are disconnected from the cloud Lightning App: example."] _ = monkeypatch_connection(monkeypatch, tmpdir, ppid=ppid_1) messages = [] - disconnect() + disconnect_app() assert "You aren't connected to any Lightning App" in messages[0] messages = [] - disconnect() + disconnect_app() assert messages == [ "You aren't connected to any Lightning App. Please use `lightning connect app_name_or_id` to connect to one." ] diff --git a/tests/tests_app/cli/test_connect_data.py b/tests/tests_app/cli/test_connect_data.py new file mode 100644 index 0000000000..aacfdb079a --- /dev/null +++ b/tests/tests_app/cli/test_connect_data.py @@ -0,0 +1,60 @@ +import sys +from unittest.mock import MagicMock + +import pytest +from lightning_cloud.openapi import ProjectIdDataConnectionsBody, V1ListMembershipsResponse, V1Membership + +from lightning.app.cli.connect import data + + +@pytest.mark.skipif(sys.platform == "win32", reason="not supported on windows yet") +def test_connect_data_no_project(monkeypatch): + + client = MagicMock() + client.projects_service_list_memberships.return_value = V1ListMembershipsResponse(memberships=[]) + monkeypatch.setattr(data, "LightningClient", MagicMock(return_value=client)) + + _error_and_exit = MagicMock() + monkeypatch.setattr(data, "_error_and_exit", _error_and_exit) + + _get_project = MagicMock() + _get_project.return_value = V1Membership(name="project-0", project_id="project-id-0") + monkeypatch.setattr(data, "_get_project", _get_project) + + data.connect_data("imagenet", region="us-east-1", source="imagenet", destination="", project_name="project-0") + + _get_project.assert_called() + + +@pytest.mark.skipif(sys.platform == "win32", reason="not supported on windows yet") +def test_connect_data(monkeypatch): + + client = MagicMock() + client.projects_service_list_memberships.return_value = V1ListMembershipsResponse( + memberships=[ + V1Membership(name="project-0", project_id="project-id-0"), + V1Membership(name="project-1", project_id="project-id-1"), + V1Membership(name="project 2", project_id="project-id-2"), + ] + ) + monkeypatch.setattr(data, "LightningClient", MagicMock(return_value=client)) + + _error_and_exit = MagicMock() + monkeypatch.setattr(data, "_error_and_exit", _error_and_exit) + data.connect_data("imagenet", region="us-east-1", source="imagenet", destination="", project_name="project-0") + + _error_and_exit.assert_called_with( + "Only public S3 folders are supported for now. Please, open a Github issue with your use case." + ) + + data.connect_data("imagenet", region="us-east-1", source="s3://imagenet", destination="", project_name="project-0") + + client.data_connection_service_create_data_connection.assert_called_with( + project_id="project-id-0", + body=ProjectIdDataConnectionsBody( + destination="", + region="us-east-1", + name="imagenet", + source="s3://imagenet", + ), + ) diff --git a/tests/tests_app/utilities/test_commands.py b/tests/tests_app/utilities/test_commands.py index b3eb2c1a21..3a96261413 100644 --- a/tests/tests_app/utilities/test_commands.py +++ b/tests/tests_app/utilities/test_commands.py @@ -10,7 +10,7 @@ from pydantic import BaseModel from lightning.app import LightningApp, LightningFlow from lightning.app.cli.commands.app_commands import _run_app_command -from lightning.app.cli.commands.connection import connect, disconnect +from lightning.app.cli.connect.app import connect_app, disconnect_app from lightning.app.core.constants import APP_SERVER_PORT from lightning.app.runners import MultiProcessRuntime from lightning.app.testing.helpers import _RunIf @@ -145,7 +145,7 @@ def test_configure_commands(monkeypatch): sleep(0.5) monkeypatch.setattr(sys, "argv", ["lightning", "user", "command", "--name=something"]) - connect("localhost") + connect_app("localhost") _run_app_command("localhost", None) sleep(0.5) state = AppState() @@ -160,5 +160,5 @@ def test_configure_commands(monkeypatch): sleep(0.1) time_left -= 0.1 assert process.exitcode == 0 - disconnect() + disconnect_app() process.kill()