Introduce Upload File endpoint (#14703)

* update

* update

* update

* update

* update

* update

* update

* update

* update

* update

Author: tchaton

src/lightning_app/CHANGELOG.md

@@ -5,7 +5,11 @@ All notable changes to this project will be documented in this file.
 The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
 
 
-## [0.6.1] - 2022-09-16
+## [0.6.1] - 2022-09-19
+
+### Added
+
+- Add support to upload files to the Drive through an asynchronous `upload_file` endpoint ([#14703](https://github.com/Lightning-AI/lightning/pull/14703))
 
 ### Changed
 

src/lightning_app/core/api.py

@@ -5,12 +5,13 @@
 import traceback
 from copy import deepcopy
 from multiprocessing import Queue
+from tempfile import TemporaryDirectory
 from threading import Event, Lock, Thread
 from typing import Dict, List, Mapping, Optional
 
 import uvicorn
 from deepdiff import DeepDiff, Delta
-from fastapi import FastAPI, HTTPException, Request, Response, status, WebSocket
+from fastapi import FastAPI, File, HTTPException, Request, Response, status, UploadFile, WebSocket
 from fastapi.middleware.cors import CORSMiddleware
 from fastapi.params import Header
 from fastapi.responses import HTMLResponse, JSONResponse
@@ -23,6 +24,7 @@
 from lightning_app.api.request_types import DeltaRequest
 from lightning_app.core.constants import FRONTEND_DIR
 from lightning_app.core.queues import RedisQueue
+from lightning_app.storage import Drive
 from lightning_app.utilities.app_helpers import InMemoryStateStore, Logger, StateStore
 from lightning_app.utilities.enum import OpenAPITags
 from lightning_app.utilities.imports import _is_redis_available, _is_starsessions_available
@@ -234,6 +236,29 @@ async def post_state(
     api_app_delta_queue.put(DeltaRequest(delta=Delta(deep_diff)))
 
 
+@fastapi_service.put("/api/v1/upload_file/{filename}")
+async def upload_file(filename: str, uploaded_file: UploadFile = File(...)):
+    with TemporaryDirectory() as tmp:
+        drive = Drive(
+            "lit://uploaded_files",
+            component_name="file_server",
+            allow_duplicates=True,
+            root_folder=tmp,
+        )
+        tmp_file = os.path.join(tmp, filename)
+
+        with open(tmp_file, "wb") as f:
+            done = False
+            while not done:
+                # Note: The 8192 number doesn't have a strong reason.
+                content = await uploaded_file.read(8192)
+                f.write(content)
+                done = content == b""
+
+        drive.put(filename)
+    return f"Successfully uploaded '{filename}' to the Drive"
+
+
 @fastapi_service.get("/healthz", status_code=200)
 async def healthz(response: Response):
     """Health check endpoint used in the cloud FastAPI servers to check the status periodically. This requires

src/lightning_app/runners/cloud.py

@@ -332,3 +332,12 @@ def _check_uploaded_folder(root: Path, repo: LocalSourceCodeDir) -> None:
             else:
                 warning_msg += "\nYou can ignore some files or folders by adding them to `.lightningignore`."
             logger.warn(warning_msg)
+
+    def _project_has_sufficient_credits(self, project: V1Membership, app: Optional[LightningApp] = None):
+        """check if user has enough credits to run the app with its hardware if app is not passed return True if
+        user has 1 or more credits."""
+        balance = project.balance
+        if balance is None:
+            balance = 0  # value is missing in some tests
+
+        return balance >= 1