[App] Add support for plugins to return actions (#16832)

Author: ethanwharris

requirements/app/base.txt

@@ -1,4 +1,4 @@
-lightning-cloud>=0.5.26
+lightning-cloud>=0.5.27
 packaging
 typing-extensions>=4.0.0, <=4.4.0
 deepdiff>=5.7.0, <6.2.4

src/lightning/app/__init__.py

@@ -33,9 +33,9 @@
 
 from lightning.app.core.app import LightningApp  # noqa: E402
 from lightning.app.core.flow import LightningFlow  # noqa: E402
-from lightning.app.core.plugin import LightningPlugin  # noqa: E402
 from lightning.app.core.work import LightningWork  # noqa: E402
 from lightning.app.perf import pdb  # noqa: E402
+from lightning.app.plugin.plugin import LightningPlugin  # noqa: E402
 from lightning.app.utilities.packaging.build_config import BuildConfig  # noqa: E402
 from lightning.app.utilities.packaging.cloud_compute import CloudCompute  # noqa: E402
 

src/lightning/app/core/__init__.py

@@ -1,6 +1,5 @@
 from lightning.app.core.app import LightningApp
 from lightning.app.core.flow import LightningFlow
-from lightning.app.core.plugin import LightningPlugin
 from lightning.app.core.work import LightningWork
 
-__all__ = ["LightningApp", "LightningFlow", "LightningWork", "LightningPlugin"]
+__all__ = ["LightningApp", "LightningFlow", "LightningWork"]

src/lightning/app/plugin/__init__.py

@@ -0,0 +1,5 @@
+import lightning.app.plugin.actions as actions
+from lightning.app.plugin.actions import NavigateTo, Toast, ToastSeverity
+from lightning.app.plugin.plugin import LightningPlugin
+
+__all__ = ["LightningPlugin", "actions", "Toast", "ToastSeverity", "NavigateTo"]

src/lightning/app/plugin/actions.py

@@ -0,0 +1,72 @@
+# Copyright The Lightning AI 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.
+from dataclasses import dataclass
+from enum import Enum
+from typing import Union
+
+from lightning_cloud.openapi.models import V1CloudSpaceAppAction, V1CloudSpaceAppActionType
+
+
+class _Action:
+    """Actions are returned by `LightningPlugin` objects to perform actions in the UI."""
+
+    def to_spec(self) -> V1CloudSpaceAppAction:
+        """Convert this action to a ``V1CloudSpaceAppAction``"""
+        raise NotImplementedError
+
+
+@dataclass
+class NavigateTo(_Action):
+    """The ``NavigateTo`` action can be used to navigate to a relative URL within the Lightning frontend.
+
+    Args:
+        url: The relative URL to navigate to. E.g. ``/<username>/<project>``.
+    """
+
+    url: str
+
+    def to_spec(self) -> V1CloudSpaceAppAction:
+        return V1CloudSpaceAppAction(
+            type=V1CloudSpaceAppActionType.NAVIGATE_TO,
+            content=self.url,
+        )
+
+
+class ToastSeverity(Enum):
+    ERROR = "error"
+    INFO = "info"
+    SUCCESS = "success"
+    WARNING = "warning"
+
+    def __str__(self) -> str:
+        return self.value
+
+
+@dataclass
+class Toast(_Action):
+    """The ``Toast`` action can be used to display a toast message to the user.
+
+    Args:
+        severity: The severity level of the toast. One of: "error", "info", "success", "warning".
+        message: The message body.
+    """
+
+    severity: Union[ToastSeverity, str]
+    message: str
+
+    def to_spec(self) -> V1CloudSpaceAppAction:
+        return V1CloudSpaceAppAction(
+            type=V1CloudSpaceAppActionType.TOAST,
+            content=f"{self.severity}:{self.message}",
+        )

src/lightning/app/core/plugin.py renamed to src/lightning/app/plugin/plugin.py

@@ -15,7 +15,7 @@
 import tarfile
 import tempfile
 from pathlib import Path
-from typing import Dict, List, Optional
+from typing import Any, Dict, List, Optional
 from urllib.parse import urlparse
 
 import requests
@@ -24,6 +24,8 @@
 from fastapi.middleware.cors import CORSMiddleware
 from pydantic import BaseModel
 
+from lightning.app.core import constants
+from lightning.app.plugin.actions import _Action
 from lightning.app.utilities.app_helpers import Logger
 from lightning.app.utilities.component import _set_flow_context
 from lightning.app.utilities.enum import AppStage
@@ -41,16 +43,20 @@ def __init__(self) -> None:
         self.cloudspace_id = ""
         self.cluster_id = ""
 
-    def run(self, *args: str, **kwargs: str) -> None:
+    def run(self, *args: str, **kwargs: str) -> Optional[List[_Action]]:
         """Override with the logic to execute on the cloudspace."""
+        raise NotImplementedError
 
-    def run_job(self, name: str, app_entrypoint: str, env_vars: Optional[Dict[str, str]] = None) -> None:
+    def run_job(self, name: str, app_entrypoint: str, env_vars: Optional[Dict[str, str]] = None) -> str:
         """Run a job in the cloudspace associated with this plugin.
 
         Args:
             name: The name of the job.
             app_entrypoint: The path of the file containing the app to run.
             env_vars: Additional env vars to set when running the app.
+
+        Returns:
+            The relative URL of the created job.
         """
         from lightning.app.runners.cloud import CloudRuntime
 
@@ -74,12 +80,14 @@ def run_job(self, name: str, app_entrypoint: str, env_vars: Optional[Dict[str, s
         # Used to indicate Lightning has been dispatched
         os.environ["LIGHTNING_DISPATCHED"] = "1"
 
-        runtime.cloudspace_dispatch(
+        url = runtime.cloudspace_dispatch(
             project_id=self.project_id,
             cloudspace_id=self.cloudspace_id,
             name=name,
             cluster_id=self.cluster_id,
         )
+        # Return a relative URL so it can be used with the NavigateTo action.
+        return url.replace(constants.get_lightning_cloud_url(), "")
 
     def _setup(
         self,
@@ -101,7 +109,7 @@ class _Run(BaseModel):
     plugin_arguments: Dict[str, str]
 
 
-def _run_plugin(run: _Run) -> List:
+def _run_plugin(run: _Run) -> Dict[str, Any]:
     """Create a run with the given name and entrypoint under the cloudspace with the given ID."""
     with tempfile.TemporaryDirectory() as tmpdir:
         download_path = os.path.join(tmpdir, "source.tar.gz")
@@ -115,6 +123,9 @@ def _run_plugin(run: _Run) -> List:
 
             response = requests.get(source_code_url)
 
+            # TODO: Backoff retry a few times in case the URL is flaky
+            response.raise_for_status()
+
             with open(download_path, "wb") as f:
                 f.write(response.content)
         except Exception as e:
@@ -152,17 +163,15 @@ def _run_plugin(run: _Run) -> List:
                 cloudspace_id=run.cloudspace_id,
                 cluster_id=run.cluster_id,
             )
-            plugin.run(**run.plugin_arguments)
+            actions = plugin.run(**run.plugin_arguments) or []
+            return {"actions": [action.to_spec().to_dict() for action in actions]}
         except Exception as e:
             raise HTTPException(
                 status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, detail=f"Error running plugin: {str(e)}."
             )
         finally:
             os.chdir(cwd)
 
-        # TODO: Return actions from the plugin here
-        return []
-
 
 def _start_plugin_server(host: str, port: int) -> None:
     """Start the plugin server which can be used to dispatch apps or run plugins."""

src/lightning/app/runners/cloud.py

@@ -196,7 +196,7 @@ def cloudspace_dispatch(
         cloudspace_id: str,
         name: str,
         cluster_id: str,
-    ):
+    ) -> str:
         """Slim dispatch for creating runs from a cloudspace. This dispatch avoids resolution of some properties
         such as the project and cluster IDs that are instead passed directly.
 
@@ -210,12 +210,15 @@ def cloudspace_dispatch(
             ApiException: If there was an issue in the backend.
             RuntimeError: If there are validation errors.
             ValueError: If there are validation errors.
+
+        Returns:
+            The URL of the created job.
         """
         # Dispatch in four phases: resolution, validation, spec creation, API transactions
         # Resolution
         root = self._resolve_root()
         repo = self._resolve_repo(root)
-        self._resolve_cloudspace(project_id, cloudspace_id)
+        project = self._resolve_project(project_id=project_id)
         existing_instances = self._resolve_run_instances_by_name(project_id, name)
         name = self._resolve_run_name(name, existing_instances)
         queue_server_type = self._resolve_queue_server_type()
@@ -240,7 +243,7 @@ def cloudspace_dispatch(
         run = self._api_create_run(project_id, cloudspace_id, run_body)
         self._api_package_and_upload_repo(repo, run)
 
-        self._api_create_run_instance(
+        run_instance = self._api_create_run_instance(
             cluster_id,
             project_id,
             name,
@@ -251,6 +254,8 @@ def cloudspace_dispatch(
             env_vars,
         )
 
+        return self._get_app_url(project, run_instance, "logs" if run.is_headless else "web-ui")
+
     def dispatch(
         self,
         name: str = "",
@@ -451,16 +456,9 @@ def _resolve_repo(
 
         return LocalSourceCodeDir(path=root, ignore_functions=ignore_functions)
 
-    def _resolve_project(self) -> V1Membership:
+    def _resolve_project(self, project_id: Optional[str] = None) -> V1Membership:
         """Determine the project to run on, choosing a default if multiple projects are found."""
-        return _get_project(self.backend.client)
-
-    def _resolve_cloudspace(self, project_id: str, cloudspace_id: str) -> V1CloudSpace:
-        """Get a cloudspace by project / cloudspace ID."""
-        return self.backend.client.cloud_space_service_get_cloud_space(
-            project_id=project_id,
-            id=cloudspace_id,
-        )
+        return _get_project(self.backend.client, project_id=project_id)
 
     def _resolve_existing_cloudspaces(self, project_id: str, cloudspace_name: str) -> List[V1CloudSpace]:
         """Lists all the cloudspaces with a name matching the provided cloudspace name."""

src/lightning/app/utilities/cloud.py

@@ -13,6 +13,7 @@
 # limitations under the License.
 
 import os
+from typing import Optional
 
 from lightning_cloud.openapi import V1Membership
 
@@ -22,10 +23,11 @@
 from lightning.app.utilities.network import LightningClient
 
 
-def _get_project(
-    client: LightningClient, project_id: str = LIGHTNING_CLOUD_PROJECT_ID, verbose: bool = True
-) -> V1Membership:
+def _get_project(client: LightningClient, project_id: Optional[str] = None, verbose: bool = True) -> V1Membership:
     """Get a project membership for the user from the backend."""
+    if project_id is None:
+        project_id = LIGHTNING_CLOUD_PROJECT_ID
+
     projects = client.projects_service_list_memberships()
     if project_id is not None:
         for membership in projects.memberships:

src/lightning/app/utilities/load_app.py

@@ -25,7 +25,7 @@
 
 if TYPE_CHECKING:
     from lightning.app import LightningApp, LightningFlow, LightningWork
-    from lightning.app.core.plugin import LightningPlugin
+    from lightning.app.plugin.plugin import LightningPlugin
 
 from lightning.app.utilities.app_helpers import _mock_missing_imports, Logger
 
@@ -85,7 +85,7 @@ def _load_objects_from_file(
 
 
 def _load_plugin_from_file(filepath: str) -> "LightningPlugin":
-    from lightning.app.core.plugin import LightningPlugin
+    from lightning.app.plugin.plugin import LightningPlugin
 
     # TODO: Plugin should be run in the context of the created main module here
     plugins, _ = _load_objects_from_file(filepath, LightningPlugin, raise_exception=True, mock_imports=False)