[App] Introduce basic auth to Lightning CLI (#16105)

* Introduce basic auth to Lightning CLI for app creation

* Parsing creds added

* Adding auth field to app instance body

* Adding tests

* Adding changelog entry

* Adding more tests


* Update runtime.py

* Setting auth on update

* Fix test

* Update lightning-cloud dep

* Apply suggestions from code review

Co-authored-by: Adrian Wälchli <[email protected]>

* Update runtime.py

* Fix for release

* Update base.txt

Co-authored-by: Jirka Borovec <[email protected]>
Co-authored-by: Adrian Wälchli <[email protected]>

Author: 3 people

src/lightning_app/CHANGELOG.md

@@ -8,7 +8,7 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
 
 ### Added
 
--
+- Added a possibility to set up basic authentication for Lightning apps ([#16105](https://github.com/Lightning-AI/lightning/pull/16105))
 
 
 ### Changed
@@ -31,7 +31,6 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
 -
 
 
-
 ## [1.8.6] - 2022-12-21
 
 ### Added

src/lightning_app/cli/lightning_cli.py

@@ -230,6 +230,7 @@ def _run_app(
     env: tuple,
     secret: tuple,
     run_app_comment_commands: bool,
+    enable_basic_auth: str,
 ) -> None:
 
     if not os.path.exists(file):
@@ -283,6 +284,7 @@ def _run_app(
         secrets=secrets,
         cluster_id=cluster_id,
         run_app_comment_commands=run_app_comment_commands,
+        enable_basic_auth=enable_basic_auth,
     )
     if runtime_type == RuntimeType.CLOUD:
         click.echo("Application is ready in the cloud")
@@ -328,6 +330,12 @@ def run() -> None:
     default=False,
     help="run environment setup commands from the app comments.",
 )
+@click.option(
+    "--enable-basic-auth",
+    type=str,
+    default="",
+    help="Enable basic authentication for the app and use credentials provided in the format username:password",
+)
 def run_app(
     file: str,
     cloud: bool,
@@ -341,6 +349,7 @@ def run_app(
     secret: tuple,
     app_args: tuple,
     run_app_comment_commands: bool,
+    enable_basic_auth: str,
 ) -> None:
     """Run an app from a file."""
     _run_app(
@@ -355,6 +364,7 @@ def run_app(
         env,
         secret,
         run_app_comment_commands,
+        enable_basic_auth,
     )
 
 

src/lightning_app/runners/cloud.py

@@ -31,6 +31,8 @@
     V1Flowserver,
     V1LightningappInstanceSpec,
     V1LightningappInstanceState,
+    V1LightningAuth,
+    V1LightningBasicAuth,
     V1LightningworkDrives,
     V1LightningworkSpec,
     V1Membership,
@@ -68,6 +70,7 @@
 from lightning_app.source_code.copytree import _filter_ignored, _parse_lightningignore
 from lightning_app.storage import Drive, Mount
 from lightning_app.utilities.app_helpers import _is_headless, Logger
+from lightning_app.utilities.auth import _credential_string_to_basic_auth_params
 from lightning_app.utilities.cloud import _get_project
 from lightning_app.utilities.dependency_caching import get_hash
 from lightning_app.utilities.load_app import load_app_from_file
@@ -211,6 +214,16 @@ def dispatch(
                 "initialize the Runtime object with `entrypoint_file` argument?"
             )
 
+        # If enable_basic_auth is set, we parse credential string and set up authentication for the app
+        auth: V1LightningAuth = None
+        if self.enable_basic_auth != "":
+            parsed_credentials = _credential_string_to_basic_auth_params(self.enable_basic_auth)
+            auth = V1LightningAuth(
+                basic=V1LightningBasicAuth(
+                    username=parsed_credentials["username"], password=parsed_credentials["password"]
+                )
+            )
+
         # Determine the root of the project: Start at the entrypoint_file and look for nearby Lightning config files,
         # going up the directory structure. The root of the project is where the Lightning config file is located.
 
@@ -296,6 +309,7 @@ def dispatch(
                 shm_size=self.app.flow_cloud_compute.shm_size,
                 preemptible=False,
             ),
+            auth=auth,
         )
 
         # if requirements file at the root of the repository is present,
@@ -473,6 +487,7 @@ def dispatch(
                             desired_state=app_release_desired_state,
                             env=v1_env_vars,
                             queue_server_type=queue_server_type,
+                            auth=app_spec.auth,
                         )
                     ),
                 )
@@ -488,6 +503,7 @@ def dispatch(
                             name=app_name,
                             env=v1_env_vars,
                             queue_server_type=queue_server_type,
+                            auth=app_spec.auth,
                         ),
                     )
                 )

src/lightning_app/runners/runtime.py

@@ -34,6 +34,7 @@ def dispatch(
     secrets: Dict[str, str] = None,
     cluster_id: str = None,
     run_app_comment_commands: bool = False,
+    enable_basic_auth: str = "",
 ) -> Optional[Any]:
     """Bootstrap and dispatch the application to the target.
 
@@ -51,6 +52,8 @@ def dispatch(
         secrets: Dict of secrets to be passed as environment variables to the app
         cluster_id: the Lightning AI cluster to run the app on. Defaults to managed Lightning AI cloud
         run_app_comment_commands: whether to parse commands from the entrypoint file and execute them before app startup
+        enable_basic_auth: whether to enable basic authentication for the app
+                           (use credentials in the format username:password as an argument)
     """
     from lightning_app.runners.runtime_type import RuntimeType
     from lightning_app.utilities.component import _set_flow_context
@@ -76,6 +79,7 @@ def dispatch(
         env_vars=env_vars,
         secrets=secrets,
         run_app_comment_commands=run_app_comment_commands,
+        enable_basic_auth=enable_basic_auth,
     )
     # Used to indicate Lightning has been dispatched
     os.environ["LIGHTNING_DISPATCHED"] = "1"
@@ -99,6 +103,7 @@ class Runtime:
     env_vars: Dict[str, str] = field(default_factory=dict)
     secrets: Dict[str, str] = field(default_factory=dict)
     run_app_comment_commands: bool = False
+    enable_basic_auth: str = ""
 
     def __post_init__(self):
         if isinstance(self.backend, str):

src/lightning_app/utilities/auth.py

@@ -0,0 +1,23 @@
+from typing import Dict
+
+
+def _credential_string_to_basic_auth_params(credential_string: str) -> Dict[str, str]:
+    """Returns the name/ID pair for each given Secret name.
+
+    Raises a `ValueError` if any of the given Secret names do not exist.
+    """
+    if credential_string.count(":") != 1:
+        raise ValueError(
+            "Credential string must follow the format username:password; "
+            + f"the provided one ('{credential_string}') does not."
+        )
+
+    username, password = credential_string.split(":")
+
+    if not username:
+        raise ValueError("Username cannot be empty.")
+
+    if not password:
+        raise ValueError("Password cannot be empty.")
+
+    return {"username": username, "password": password}

tests/tests_app/cli/test_run_app.py

@@ -69,6 +69,7 @@ def test_lightning_run_cluster_without_cloud(monkeypatch):
             env=("FOO=bar",),
             secret=(),
             run_app_comment_commands=False,
+            enable_basic_auth="",
         )
 
 
@@ -96,6 +97,7 @@ def test_lightning_run_app_cloud(mock_dispatch: mock.MagicMock, open_ui, caplog,
             env=("FOO=bar",),
             secret=("BAR=my-secret",),
             run_app_comment_commands=False,
+            enable_basic_auth="",
         )
     # capture logs.
     # TODO(yurij): refactor the test, check if the actual HTTP request is being sent and that the proper admin
@@ -112,6 +114,7 @@ def test_lightning_run_app_cloud(mock_dispatch: mock.MagicMock, open_ui, caplog,
         secrets={"BAR": "my-secret"},
         cluster_id="",
         run_app_comment_commands=False,
+        enable_basic_auth="",
     )
 
 
@@ -139,6 +142,7 @@ def test_lightning_run_app_cloud_with_run_app_commands(mock_dispatch: mock.Magic
             env=("FOO=bar",),
             secret=("BAR=my-secret",),
             run_app_comment_commands=True,
+            enable_basic_auth="",
         )
     # capture logs.
     # TODO(yurij): refactor the test, check if the actual HTTP request is being sent and that the proper admin
@@ -155,6 +159,7 @@ def test_lightning_run_app_cloud_with_run_app_commands(mock_dispatch: mock.Magic
         secrets={"BAR": "my-secret"},
         cluster_id="",
         run_app_comment_commands=True,
+        enable_basic_auth="",
     )
 
 
@@ -175,4 +180,45 @@ def test_lightning_run_app_secrets(monkeypatch):
             env=(),
             secret=("FOO=my-secret"),
             run_app_comment_commands=False,
+            enable_basic_auth="",
         )
+
+
+@mock.patch.dict(os.environ, {"LIGHTNING_CLOUD_URL": "https://beta.lightning.ai"})
+@mock.patch("lightning_app.cli.lightning_cli.dispatch")
+def test_lightning_run_app_enable_basic_auth_passed(mock_dispatch: mock.MagicMock, caplog, monkeypatch):
+    """This test just validates the command has ran properly when --enable-basic-auth argument is passed.
+
+    It checks the call to `dispatch` for the right arguments.
+    """
+    monkeypatch.setattr("lightning_app.runners.cloud.logger", logging.getLogger())
+
+    with caplog.at_level(logging.INFO):
+        _run_app(
+            file=os.path.join(_PROJECT_ROOT, "tests/tests_app/core/scripts/app_metadata.py"),
+            cloud=True,
+            cluster_id="",
+            without_server=False,
+            name="",
+            blocking=False,
+            open_ui=False,
+            no_cache=True,
+            env=("FOO=bar",),
+            secret=("BAR=my-secret",),
+            run_app_comment_commands=False,
+            enable_basic_auth="username:password",
+        )
+    mock_dispatch.assert_called_with(
+        Path(os.path.join(_PROJECT_ROOT, "tests/tests_app/core/scripts/app_metadata.py")),
+        RuntimeType.CLOUD,
+        start_server=True,
+        blocking=False,
+        open_ui=False,
+        name="",
+        no_cache=True,
+        env_vars={"FOO": "bar"},
+        secrets={"BAR": "my-secret"},
+        cluster_id="",
+        run_app_comment_commands=False,
+        enable_basic_auth="username:password",
+    )

tests/tests_app/runners/test_cloud.py

@@ -27,6 +27,8 @@
     V1GetClusterResponse,
     V1LightningappInstanceState,
     V1LightningappRelease,
+    V1LightningAuth,
+    V1LightningBasicAuth,
     V1LightningworkDrives,
     V1LightningworkSpec,
     V1ListClustersResponse,
@@ -459,6 +461,65 @@ def test_requirements_file(self, monkeypatch):
             project_id="test-project-id", app_id=mock.ANY, body=body
         )
 
+    @mock.patch("lightning_app.runners.backends.cloud.LightningClient", mock.MagicMock())
+    def test_basic_auth_enabled(self, monkeypatch):
+        mock_client = mock.MagicMock()
+        mock_client.projects_service_list_memberships.return_value = V1ListMembershipsResponse(
+            memberships=[V1Membership(name="test-project", project_id="test-project-id")]
+        )
+        mock_client.lightningapp_instance_service_list_lightningapp_instances.return_value = (
+            V1ListLightningappInstancesResponse(lightningapps=[])
+        )
+        mock_client.lightningapp_v2_service_create_lightningapp_release.return_value = V1LightningappRelease()
+        mock_client.cluster_service_list_clusters.return_value = V1ListClustersResponse([Externalv1Cluster(id="test")])
+        cloud_backend = mock.MagicMock()
+        cloud_backend.client = mock_client
+        monkeypatch.setattr(backends, "CloudBackend", mock.MagicMock(return_value=cloud_backend))
+        monkeypatch.setattr(cloud, "LocalSourceCodeDir", mock.MagicMock())
+        monkeypatch.setattr(cloud, "_prepare_lightning_wheels_and_requirements", mock.MagicMock())
+        app = mock.MagicMock()
+        app.is_headless = False
+        app.flows = []
+        app.frontend = {}
+        cloud_runtime = cloud.CloudRuntime(app=app, entrypoint_file="entrypoint.py")
+        cloud_runtime._check_uploaded_folder = mock.MagicMock()
+        # Set cloud_runtime.enable_basic_auth to be not empty:
+        cloud_runtime.enable_basic_auth = "username:password"
+        monkeypatch.setattr(Path, "is_file", lambda *args, **kwargs: False)
+        monkeypatch.setattr(cloud, "Path", Path)
+
+        cloud_runtime.dispatch()
+        mock_client = cloud_runtime.backend.client
+
+        body = Body8(
+            app_entrypoint_file=mock.ANY,
+            enable_app_server=True,
+            is_headless=False,
+            flow_servers=[],
+            image_spec=None,
+            works=[],
+            local_source=True,
+            dependency_cache_key=mock.ANY,
+            user_requested_flow_compute_config=mock.ANY,
+        )
+
+        mock_client.lightningapp_v2_service_create_lightningapp_release.assert_called_once_with(
+            project_id="test-project-id", app_id=mock.ANY, body=body
+        )
+
+        mock_client.lightningapp_v2_service_create_lightningapp_release_instance.assert_called_once_with(
+            project_id="test-project-id",
+            app_id=mock.ANY,
+            id=mock.ANY,
+            body=Body9(
+                desired_state=mock.ANY,
+                name=mock.ANY,
+                env=mock.ANY,
+                queue_server_type=mock.ANY,
+                auth=V1LightningAuth(basic=V1LightningBasicAuth(username="username", password="password")),
+            ),
+        )
+
     @mock.patch("lightning_app.runners.backends.cloud.LightningClient", mock.MagicMock())
     def test_no_cache(self, monkeypatch):
         mock_client = mock.MagicMock()

tests/tests_app/utilities/test_auth.py

@@ -0,0 +1,26 @@
+from typing import Dict
+
+import pytest
+
+from lightning_app.utilities.auth import _credential_string_to_basic_auth_params
+
+
+@pytest.mark.parametrize(
+    "credential_string, expected_parsed, exception_message",
+    [
+        ("", None, "Credential string must follow the format username:password; the provided one ('') does not."),
+        (":", None, "Username cannot be empty."),
+        (":pass", None, "Username cannot be empty."),
+        ("user:", None, "Password cannot be empty."),
+        ("user:pass", {"username": "user", "password": "pass"}, ""),
+    ],
+)
+def test__credential_string_to_basic_auth_params(
+    credential_string: str, expected_parsed: Dict[str, str], exception_message: str
+):
+    if expected_parsed is not None:
+        assert _credential_string_to_basic_auth_params(credential_string) == expected_parsed
+    else:
+        with pytest.raises(ValueError) as exception:
+            _credential_string_to_basic_auth_params(credential_string)
+        assert exception_message == str(exception.value)