[App] Enable uploading files to a project. (#16631)

* update

* update

* update

* update

* update

* update

* update

* update

* update

* update

* update

* update

* update

* update

---------

Co-authored-by: thomas <[email protected]>

(cherry picked from commit 55a0caa)

Author: tchaton

requirements/app/base.txt

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

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 `cp` (upload) at project level ([#16631](https://github.com/Lightning-AI/lightning/pull/16631))
 - 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
@@ -41,6 +42,7 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
 
 - Fixed the Drive root_folder not parsed properly ([#16454](https://github.com/Lightning-AI/lightning/pull/16454))
 
+- Fixed malformed path when downloading files using `lightning cp` ([#16626](https://github.com/Lightning-AI/lightning/pull/16626))
 
 ## [1.9.0] - 2023-01-17
 

src/lightning_app/cli/commands/cp.py

@@ -22,9 +22,8 @@
 
 import click
 import requests
-import rich
 import urllib3
-from lightning_cloud.openapi import Externalv1LightningappInstance, IdArtifactsBody, V1CloudSpace
+from lightning_cloud.openapi import Externalv1LightningappInstance, ProjectIdStorageBody, V1CloudSpace
 from rich.live import Live
 from rich.progress import BarColumn, DownloadColumn, Progress, Task, TextColumn
 from rich.spinner import Spinner
@@ -34,6 +33,7 @@
 from lightning_app.cli.commands.pwd import _pwd
 from lightning_app.source_code import FileUploader
 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
 
 logger = Logger(__name__)
@@ -55,7 +55,7 @@ def cp(src_path: str, dst_path: str, r: bool = False, recursive: bool = False) -
         pwd = _pwd()
 
         if pwd == "/" or len(pwd.split("/")) == 1:
-            return _error_and_exit("Uploading files at the project level isn't supported yet.")
+            return _error_and_exit("Uploading files at the project level isn't allowed yet.")
 
         client = LightningClient()
 
@@ -74,10 +74,18 @@ def cp(src_path: str, dst_path: str, r: bool = False, recursive: bool = False) -
 
 
 def _upload_files(live, client: LightningClient, local_src: str, remote_dst: str, pwd: str) -> str:
+    remote_splits = [split for split in remote_dst.split("/") if split != ""]
+    remote_dst = os.path.join(*remote_splits)
+
     if not os.path.exists(local_src):
         return _error_and_exit(f"The provided source path {local_src} doesn't exist.")
 
-    project_id, app_id = _get_project_app_ids(pwd)
+    lit_resource = None
+
+    if len(remote_splits) > 1:
+        project_id, lit_resource = _get_project_id_and_resource(pwd)
+    else:
+        project_id = _get_project_id_from_name(remote_dst)
 
     local_src = Path(local_src).resolve()
     upload_paths = []
@@ -91,21 +99,34 @@ def _upload_files(live, client: LightningClient, local_src: str, remote_dst: str
 
     upload_urls = []
 
+    clusters = client.projects_service_list_project_cluster_bindings(project_id)
+
     for upload_path in upload_paths:
-        filename = str(upload_path).replace(str(os.getcwd()), "")[1:]
-        response = client.lightningapp_instance_service_upload_lightningapp_instance_artifact(
-            project_id=project_id,
-            id=app_id,
-            body=IdArtifactsBody(filename),
-            async_req=True,
-        )
-        upload_urls.append(response)
+        for cluster in clusters.clusters:
+            filename = str(upload_path).replace(str(os.getcwd()), "")[1:]
+            if lit_resource:
+                filename = _get_prefix(os.path.join(remote_dst, filename), lit_resource)
+            else:
+                filename = "/" + filename
+
+            response = client.lightningapp_instance_service_upload_project_artifact(
+                project_id=project_id,
+                body=ProjectIdStorageBody(cluster_id=cluster.cluster_id, filename=filename),
+                async_req=True,
+            )
+            upload_urls.append(response)
+
+    upload_urls = [upload_url.get().upload_url for upload_url in upload_urls]
 
     live.stop()
 
+    if not upload_paths:
+        print("There were no files to upload.")
+        return
+
     progress = _get_progress_bar()
 
-    total_size = sum([Path(path).stat().st_size for path in upload_paths])
+    total_size = sum([Path(path).stat().st_size for path in upload_paths]) // max(len(clusters.clusters), 1)
     task_id = progress.add_task("upload", filename="", total=total_size)
 
     progress.start()
@@ -126,7 +147,7 @@ def _upload_files(live, client: LightningClient, local_src: str, remote_dst: str
 def _upload(source_file: str, presigned_url: ApplyResult, progress: Progress, task_id: Task) -> Optional[Exception]:
     source_file = Path(source_file)
     file_uploader = FileUploader(
-        presigned_url.get().upload_url,
+        presigned_url,
         source_file,
         total_size=None,
         name=str(source_file),
@@ -143,13 +164,13 @@ def _download_files(live, client, remote_src: str, local_dst: str, pwd: str):
     download_urls = []
     total_size = []
 
-    prefix = _get_prefix("/".join(pwd.split("/")[3:]), lit_resource)
+    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)
-        download_paths.append(Path(path).resolve())
+        download_paths.append(path)
         download_urls.append(artifact.url)
         total_size.append(int(artifact.size_bytes))
 
@@ -182,14 +203,17 @@ def _download_file(path: str, url: str, progress: Progress, task_id: Task) -> No
     # Disable warning about making an insecure request
     urllib3.disable_warnings(urllib3.exceptions.InsecureRequestWarning)
 
-    request = requests.get(url, stream=True, verify=False)
+    try:
+        request = requests.get(url, stream=True, verify=False)
 
-    chunk_size = 1024
+        chunk_size = 1024
 
-    with open(path, "wb") as fp:
-        for chunk in request.iter_content(chunk_size=chunk_size):
-            fp.write(chunk)  # type: ignore
-            progress.update(task_id, advance=len(chunk))
+        with open(path, "wb") as fp:
+            for chunk in request.iter_content(chunk_size=chunk_size):
+                fp.write(chunk)  # type: ignore
+                progress.update(task_id, advance=len(chunk))
+    except ConnectionError:
+        pass
 
 
 def _sanitize_path(path: str, pwd: str) -> Tuple[str, bool]:
@@ -211,29 +235,6 @@ def _remove_remote(path: str) -> str:
     return path.replace("r:", "").replace("remote:", "")
 
 
-def _error_and_exit(msg: str) -> str:
-    rich.print(f"[red]ERROR[/red]: {msg}")
-    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
-    project_name, app_name, *_ = pwd.split("/")[1:3]
-    client = LightningClient()
-    projects = client.projects_service_list_memberships()
-    project_id = [project.project_id for project in projects.memberships if project.name == project_name][0]
-    client = LightningClient()
-    lit_apps = client.lightningapp_instance_service_list_lightningapp_instances(project_id=project_id).lightningapps
-    lit_apps = [lit_app for lit_app in lit_apps if lit_app.name == app_name]
-    if len(lit_apps) != 1:
-        print(f"ERROR: There isn't any Lightning App matching the name {app_name}.")
-        sys.exit(0)
-    lit_app = lit_apps[0]
-    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
@@ -263,6 +264,13 @@ def _get_project_id_and_resource(pwd: str) -> Tuple[str, Union[Externalv1Lightni
     return project_id, lit_ressources[0]
 
 
+def _get_project_id_from_name(project_name: str) -> str:
+    # 1. Collect the projects of the user
+    client = LightningClient()
+    projects = client.projects_service_list_memberships()
+    return [project.project_id for project in projects.memberships if project.name == project_name][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 fastapi import HTTPException
 from lightning_cloud.openapi import Externalv1LightningappInstance
 from rich.console import Console
 from rich.live import Live
@@ -27,6 +28,7 @@
 from lightning_app.cli.commands.cd import _CD_FILE
 from lightning_app.cli.commands.connection 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
 
 _FOLDER_COLOR = "sky_blue1"
@@ -71,7 +73,15 @@ def ls(path: Optional[str] = None) -> List[str]:
         # TODO: Add support for CloudSpaces, etc..
         splits = root.split("/")[1:]
 
-        project_id = [project.project_id for project in projects.memberships if project.name == splits[0]][0]
+        project = [project for project in projects.memberships if project.name == splits[0]]
+
+        # This happens if the user changes cluster and the project doesn't exit.
+        if len(project) == 0:
+            return _error_and_exit(
+                f"There isn't any Lightning Project matching the name {splits[0]}." " HINT: Use `lightning cd`."
+            )
+
+        project_id = project[0].project_id
 
         lit_apps = client.lightningapp_instance_service_list_lightningapp_instances(project_id=project_id).lightningapps
 
@@ -91,9 +101,7 @@ def ls(path: Optional[str] = None) -> List[str]:
             lit_ressources = [lit_resource for lit_resource in lit_apps if lit_resource.name == splits[1]]
 
             if len(lit_ressources) == 0:
-
-                print(f"ERROR: There isn't any Lightning Ressource matching the name {splits[1]}.")
-                sys.exit(0)
+                _error_and_exit(f"There isn't any Lightning Ressource matching the name {splits[1]}.")
 
         lit_resource = lit_ressources[0]
 
@@ -212,32 +220,44 @@ def _collect_artifacts(
         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,
+        try:
+            response = client.lightningapp_instance_service_list_project_artifacts(
                 project_id,
                 prefix=prefix,
                 cluster_id=cluster_id,
-                page_token=response.next_page_token,
-                tokens=tokens,
+                page_token=page_token,
+                include_download_url=include_download_url,
+                page_size=str(page_size),
             )
+            if response:
+                for artifact in response.artifacts:
+                    if ".lightning-app-sync" in artifact.filename:
+                        continue
+                    yield artifact
+
+                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,
+                    )
+        except HTTPException:
+            # Note: This is triggered when the request is wrong.
+            # This is currently happening due to looping through the user clusters.
+            pass
 
 
 def _add_resource_prefix(prefix: str, resource_path: str):
     if resource_path in prefix:
         return prefix
-    return "/" + os.path.join(resource_path, prefix)
+    prefix = os.path.join(resource_path, prefix)
+    if not prefix.startswith("/"):
+        prefix = "/" + prefix
+    return prefix
 
 
 def _get_prefix(prefix: str, lit_resource) -> str:

src/lightning_app/storage/filesystem.py

@@ -50,8 +50,13 @@ def put(self, src_path: str, dst_path: str, put_fn: Callable = _copy_files) -> N
         if not dst_path.startswith("/"):
             raise Exception(f"The provided destination {dst_path} needs to start with `/`.")
 
+        if dst_path == "/":
+            dst_path = os.path.join(self._root, os.path.basename(src_path))
+        else:
+            dst_path = os.path.join(self._root, dst_path[1:])
+
         src = Path(src_path).resolve()
-        dst = Path(os.path.join(self._root, dst_path[1:])).resolve()
+        dst = Path(dst_path).resolve()
 
         return put_fn(src, dst, fs=self._fs)
 

src/lightning_app/utilities/cli_helpers.py

@@ -25,6 +25,7 @@
 import click
 import packaging
 import requests
+import rich
 
 from lightning_app import __package_name__, __version__
 from lightning_app.core.constants import APP_SERVER_PORT
@@ -349,3 +350,8 @@ def _check_environment_and_redirect():
 
         _redirect_command(env_executable)
     return
+
+
+def _error_and_exit(msg: str) -> str:
+    rich.print(f"[red]ERROR[/red]: {msg}")
+    sys.exit(0)

src/lightning_app/utilities/network.py

@@ -12,6 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import ast
 import socket
 import time
 from functools import wraps
@@ -21,6 +22,7 @@
 import lightning_cloud
 import requests
 import urllib3
+from fastapi import HTTPException
 from lightning_cloud.rest_client import create_swagger_client, GridRestClient
 from requests import Session
 from requests.adapters import HTTPAdapter
@@ -137,6 +139,8 @@ def wrapped(*args: Any, **kwargs: Any) -> Any:
             try:
                 return func(self, *args, **kwargs)
             except lightning_cloud.openapi.rest.ApiException as e:
+                if e.status == 500:
+                    raise HTTPException(status_code=500, detail=ast.literal_eval(e.body.decode("utf-8"))["message"])
                 # retry if the control plane fails with all errors except 4xx but not 408 - (Request Timeout)
                 if e.status == 408 or e.status == 409 or not str(e.status).startswith("4"):
                     consecutive_errors += 1