[App] Enable listing at project level 2/n (#16622)

Co-authored-by: Ethan Harris <[email protected]>
Co-authored-by: thomas <[email protected]>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
(cherry picked from commit 03bb29c)

Author: 4 people

src/lightning_app/CHANGELOG.md

@@ -13,6 +13,7 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
 
 - Added FileSystem abstraction to simply manipulation of files ([#16581](https://github.com/Lightning-AI/lightning/pull/16581))
 
+- Enabled `ls` and `cp` (download) at project level ([#16622](https://github.com/Lightning-AI/lightning/pull/16622))
 - Added Storage Commands ([#16606](https://github.com/Lightning-AI/lightning/pull/16606))
  * `ls`: List files from your Cloud Platform Filesystem
  * `cd`: Change the current directory within your Cloud Platform filesystem (terminal session based)

src/lightning_app/cli/commands/cd.py

@@ -37,8 +37,6 @@ def cd(path: Optional[Union[Tuple[str], str]]) -> None:
 
         root = "/"
 
-        live.stop()
-
         if isinstance(path, Tuple) and len(path) > 0:
             path = " ".join(path)
 

src/lightning_app/cli/commands/cp.py

@@ -18,14 +18,13 @@
 from functools import partial
 from multiprocessing.pool import ApplyResult
 from pathlib import Path
-from time import sleep
-from typing import Optional, Tuple
+from typing import Optional, Tuple, Union
 
 import click
 import requests
 import rich
 import urllib3
-from lightning_cloud.openapi import IdArtifactsBody
+from lightning_cloud.openapi import Externalv1LightningappInstance, IdArtifactsBody, V1CloudSpace
 from rich.live import Live
 from rich.progress import BarColumn, DownloadColumn, Progress, Task, TextColumn
 from rich.spinner import Spinner
@@ -60,8 +59,8 @@ def cp(src_path: str, dst_path: str, r: bool = False, recursive: bool = False) -
 
         client = LightningClient()
 
-        src_path, src_remote = _sanetize_path(src_path, pwd)
-        dst_path, dst_remote = _sanetize_path(dst_path, pwd)
+        src_path, src_remote = _sanitize_path(src_path, pwd)
+        dst_path, dst_remote = _sanitize_path(dst_path, pwd)
 
         if src_remote and dst_remote:
             return _error_and_exit("Moving files remotely isn't supported yet. Please, open a Github issue.")
@@ -104,9 +103,6 @@ def _upload_files(live, client: LightningClient, local_src: str, remote_dst: str
 
     live.stop()
 
-    # Sleep to avoid rich live collision.
-    sleep(1)
-
     progress = _get_progress_bar()
 
     total_size = sum([Path(path).stat().st_size for path in upload_paths])
@@ -141,14 +137,15 @@ def _upload(source_file: str, presigned_url: ApplyResult, progress: Progress, ta
 
 
 def _download_files(live, client, remote_src: str, local_dst: str, pwd: str):
-    project_id, app_id = _get_project_app_ids(pwd)
+    project_id, lit_resource = _get_project_id_and_resource(pwd)
 
     download_paths = []
     download_urls = []
     total_size = []
 
-    response = client.lightningapp_instance_service_list_lightningapp_instance_artifacts(project_id, app_id)
-    for artifact in response.artifacts:
+    prefix = _get_prefix("/".join(pwd.split("/")[3:]), lit_resource)
+
+    for artifact in _collect_artifacts(client, project_id, prefix, include_download_url=True):
         path = os.path.join(local_dst, artifact.filename.replace(remote_src, ""))
         path = Path(path).resolve()
         os.makedirs(path.parent, exist_ok=True)
@@ -158,14 +155,15 @@ def _download_files(live, client, remote_src: str, local_dst: str, pwd: str):
 
     live.stop()
 
-    # Sleep to avoid rich live collision.
-    sleep(1)
+    if not download_paths:
+        print("There were no files to download.")
+        return
 
     progress = progress = _get_progress_bar()
 
     progress.start()
 
-    task_id = progress.add_task("download", filename=path, total=sum(total_size))
+    task_id = progress.add_task("download", filename="", total=sum(total_size))
 
     _download_file_fn = partial(_download_file, progress=progress, task_id=task_id)
 
@@ -194,7 +192,7 @@ def _download_file(path: str, url: str, progress: Progress, task_id: Task) -> No
             progress.update(task_id, advance=len(chunk))
 
 
-def _sanetize_path(path: str, pwd: str) -> Tuple[str, bool]:
+def _sanitize_path(path: str, pwd: str) -> Tuple[str, bool]:
     is_remote = _is_remote(path)
     if is_remote:
         path = _remove_remote(path)
@@ -218,6 +216,7 @@ def _error_and_exit(msg: str) -> str:
     sys.exit(0)
 
 
+# TODO: To be removed when upload is supported for CloudSpaces.
 def _get_project_app_ids(pwd: str) -> Tuple[str, str]:
     """Convert a root path to a project id and app id."""
     # TODO: Handle project level
@@ -235,6 +234,35 @@ def _get_project_app_ids(pwd: str) -> Tuple[str, str]:
     return project_id, lit_app.id
 
 
+def _get_project_id_and_resource(pwd: str) -> Tuple[str, Union[Externalv1LightningappInstance, V1CloudSpace]]:
+    """Convert a root path to a project id and app id."""
+    # TODO: Handle project level
+    project_name, resource_name, *_ = pwd.split("/")[1:3]
+
+    # 1. Collect the projects of the user
+    client = LightningClient()
+    projects = client.projects_service_list_memberships()
+    project_id = [project.project_id for project in projects.memberships if project.name == project_name][0]
+
+    # 2. Collect resources
+    lit_apps = client.lightningapp_instance_service_list_lightningapp_instances(project_id=project_id).lightningapps
+
+    lit_cloud_spaces = client.cloud_space_service_list_cloud_spaces(project_id=project_id).cloudspaces
+
+    lit_ressources = [lit_resource for lit_resource in lit_cloud_spaces if lit_resource.name == resource_name]
+
+    if len(lit_ressources) == 0:
+
+        lit_ressources = [lit_resource for lit_resource in lit_apps if lit_resource.name == resource_name]
+
+        if len(lit_ressources) == 0:
+
+            print(f"ERROR: There isn't any Lightning Ressource matching the name {resource_name}.")
+            sys.exit(0)
+
+    return project_id, lit_ressources[0]
+
+
 def _get_progress_bar():
     return Progress(
         TextColumn("[bold blue]{task.description}", justify="left"),

src/lightning_app/cli/commands/ls.py

@@ -18,6 +18,7 @@
 
 import click
 import rich
+from lightning_cloud.openapi import Externalv1LightningappInstance
 from rich.console import Console
 from rich.live import Live
 from rich.spinner import Spinner
@@ -44,9 +45,7 @@ def ls(path: Optional[str] = None) -> List[str]:
 
     root = "/"
 
-    with Live(Spinner("point", text=Text("pending...", style="white")), transient=True) as live:
-
-        live.stop()
+    with Live(Spinner("point", text=Text("pending...", style="white")), transient=True):
 
         if not os.path.exists(_LIGHTNING_CONNECTION_FOLDER):
             os.makedirs(_LIGHTNING_CONNECTION_FOLDER)
@@ -76,45 +75,74 @@ def ls(path: Optional[str] = None) -> List[str]:
 
         lit_apps = client.lightningapp_instance_service_list_lightningapp_instances(project_id=project_id).lightningapps
 
+        lit_cloud_spaces = client.cloud_space_service_list_cloud_spaces(project_id=project_id).cloudspaces
+
         if len(splits) == 1:
-            app_names = sorted([lit_app.name for lit_app in lit_apps])
-            _print_names_with_colors(app_names, [_FOLDER_COLOR] * len(app_names))
-            return app_names
+            apps = [lit_app.name for lit_app in lit_apps]
+            cloud_spaces = [lit_cloud_space.name for lit_cloud_space in lit_cloud_spaces]
+            ressource_names = sorted(set(cloud_spaces + apps))
+            _print_names_with_colors(ressource_names, [_FOLDER_COLOR] * len(ressource_names))
+            return ressource_names
+
+        lit_ressources = [lit_resource for lit_resource in lit_cloud_spaces if lit_resource.name == splits[1]]
+
+        if len(lit_ressources) == 0:
 
-        lit_apps = [lit_app for lit_app in lit_apps if lit_app.name == splits[1]]
+            lit_ressources = [lit_resource for lit_resource in lit_apps if lit_resource.name == splits[1]]
 
-        if len(lit_apps) != 1:
-            print(f"ERROR: There isn't any Lightning App matching the name {splits[1]}.")
-            sys.exit(0)
+            if len(lit_ressources) == 0:
 
-        lit_app = lit_apps[0]
+                print(f"ERROR: There isn't any Lightning Ressource matching the name {splits[1]}.")
+                sys.exit(0)
+
+        lit_resource = lit_ressources[0]
+
+        app_paths = []
+        app_colors = []
+
+        cloud_spaces_paths = []
+        cloud_spaces_colors = []
 
-        paths = []
-        colors = []
         depth = len(splits)
-        subpath = "/".join(splits[2:])
-        # TODO: Replace with project level endpoints
-        for artifact in _collect_artifacts(client, project_id, lit_app.id):
-            path = os.path.join(project_id, lit_app.name, artifact.filename)
+
+        prefix = "/".join(splits[2:])
+        prefix = _get_prefix(prefix, lit_resource)
+
+        for artifact in _collect_artifacts(client=client, project_id=project_id, prefix=prefix):
+
+            if str(artifact.filename).startswith("/"):
+                artifact.filename = artifact.filename[1:]
+
+            path = os.path.join(project_id, prefix[1:], artifact.filename)
+
             artifact_splits = path.split("/")
 
-            if len(artifact_splits) < depth + 1:
+            if len(artifact_splits) <= depth + 1:
                 continue
 
-            if not str(artifact.filename).startswith(subpath):
-                continue
+            path = artifact_splits[depth + 1]
 
-            path = artifact_splits[depth]
+            paths = app_paths if isinstance(lit_resource, Externalv1LightningappInstance) else cloud_spaces_paths
+            colors = app_colors if isinstance(lit_resource, Externalv1LightningappInstance) else cloud_spaces_colors
 
             if path not in paths:
                 paths.append(path)
 
                 # display files otherwise folders
                 colors.append(_FILE_COLOR if len(artifact_splits) == depth + 1 else _FOLDER_COLOR)
 
-    _print_names_with_colors(paths, colors)
+    if app_paths and cloud_spaces_paths:
+        if app_paths:
+            rich.print("Lightning App")
+            _print_names_with_colors(app_paths, app_colors)
+
+        if cloud_spaces_paths:
+            rich.print("Lightning CloudSpaces")
+            _print_names_with_colors(cloud_spaces_paths, cloud_spaces_colors)
+    else:
+        _print_names_with_colors(app_paths + cloud_spaces_paths, app_colors + cloud_spaces_colors)
 
-    return paths
+    return app_paths + cloud_spaces_paths
 
 
 def _add_colors(filename: str, color: Optional[str] = None) -> str:
@@ -156,21 +184,64 @@ def _print_names_with_colors(names: List[str], colors: List[str], padding: int =
 def _collect_artifacts(
     client: LightningClient,
     project_id: str,
-    app_id: str,
+    prefix: str = "",
     page_token: Optional[str] = "",
+    cluster_id: Optional[str] = None,
+    page_size: int = 100_000,
     tokens=None,
+    include_download_url: bool = False,
 ) -> Generator:
     if tokens is None:
         tokens = []
 
-    if page_token in tokens:
-        return
-
-    response = client.lightningapp_instance_service_list_lightningapp_instance_artifacts(
-        project_id, app_id, page_token=page_token
-    )
-    yield from response.artifacts
-
-    if response.next_page_token != "":
-        tokens.append(page_token)
-        yield from _collect_artifacts(client, project_id, app_id, page_token=response.next_page_token, tokens=tokens)
+    if cluster_id is None:
+        clusters = client.projects_service_list_project_cluster_bindings(project_id)
+        for cluster in clusters.clusters:
+            yield from _collect_artifacts(
+                client,
+                project_id,
+                prefix=prefix,
+                cluster_id=cluster.cluster_id,
+                page_token=page_token,
+                tokens=tokens,
+                page_size=page_size,
+                include_download_url=include_download_url,
+            )
+    else:
+
+        if page_token in tokens:
+            return
+
+        response = client.lightningapp_instance_service_list_project_artifacts(
+            project_id,
+            prefix=prefix,
+            cluster_id=cluster_id,
+            page_token=page_token,
+            include_download_url=include_download_url,
+            page_size=str(page_size),
+        )
+        yield from response.artifacts
+
+        if response.next_page_token != "":
+            tokens.append(page_token)
+            yield from _collect_artifacts(
+                client,
+                project_id,
+                prefix=prefix,
+                cluster_id=cluster_id,
+                page_token=response.next_page_token,
+                tokens=tokens,
+            )
+
+
+def _add_resource_prefix(prefix: str, resource_path: str):
+    if resource_path in prefix:
+        return prefix
+    return "/" + os.path.join(resource_path, prefix)
+
+
+def _get_prefix(prefix: str, lit_resource) -> str:
+    if isinstance(lit_resource, Externalv1LightningappInstance):
+        return _add_resource_prefix(prefix, f"lightningapps/{lit_resource.id}")
+
+    return _add_resource_prefix(prefix, f"cloudspaces/{lit_resource.id}")

tests/tests_app/cli/test_cp.py

@@ -71,6 +71,10 @@ def test_cp_cloud_to_local(tmpdir, monkeypatch):
         memberships=[V1Membership(name="project-0")]
     )
 
+    clusters = MagicMock()
+    clusters.clusters = [MagicMock()]
+    client.projects_service_list_project_cluster_bindings.return_value = clusters
+
     client.lightningapp_instance_service_list_lightningapp_instances.return_value = V1ListLightningappInstancesResponse(
         lightningapps=[
             Externalv1LightningappInstance(
@@ -80,7 +84,7 @@ def test_cp_cloud_to_local(tmpdir, monkeypatch):
         ]
     )
 
-    client.lightningapp_instance_service_list_lightningapp_instance_artifacts.return_value = (
+    client.lightningapp_instance_service_list_project_artifacts.return_value = (
         V1ListLightningappInstanceArtifactsResponse(
             artifacts=[
                 V1LightningappInstanceArtifact(