Introduce Upload File endpoint

src/lightning_app/CHANGELOG.md

@@ -12,6 +12,9 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
 - Add `load_state_dict` and `state_dict` ([#14100](https://github.com/Lightning-AI/lightning/pull/14100))
 
 
+- Add support to upload files to the Drive through an asynchronous `upload_file` endpoint ([#14703](https://github.com/Lightning-AI/lightning/pull/14703))
+
+
 ### Changed
 
 - Application storage prefix moved from `app_id` to `project_id/app_id` ([#14583](https://github.com/Lightning-AI/lightning/pull/14583))

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 ENABLE_STATE_WEBSOCKET, 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: Optional[str] = None, uploaded_file: UploadFile = File(...)):
+    try:
+        filename = filename if filename else uploaded_file.filename
+        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:
+                while content := await uploaded_file.read(10240):
+                    f.write(content)
+
+            drive.put(filename)
+    finally:
+        uploaded_file.file.close()
+    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

@@ -342,7 +342,7 @@ 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.warning(warning_msg)
+            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

tests/tests_app/core/test_lightning_api.py

@@ -25,6 +25,7 @@
 from lightning_app.core.constants import APP_SERVER_PORT
 from lightning_app.runners import MultiProcessRuntime, SingleProcessRuntime
 from lightning_app.storage.drive import Drive
+from lightning_app.storage.path import shared_storage_path
 from lightning_app.testing.helpers import MockQueue
 from lightning_app.utilities.component import _set_frontend_context, _set_work_context
 from lightning_app.utilities.enum import AppStage
@@ -429,7 +430,7 @@ def target():
     MultiProcessRuntime(app).dispatch()
 
 
-def test_configure_api():
+def test_configure_api_and_upload_files():
 
     process = Process(target=target)
     process.start()
@@ -446,6 +447,27 @@ def test_configure_api():
         f"http://localhost:{APP_SERVER_PORT}/api/v1/request", data=InputRequestModel(name="hello").json()
     )
     assert response.json() == {"name": "hello", "counter": 1}
+
+    with open(__file__, "rb") as f:
+        response = requests.put(
+            f"http://localhost:{APP_SERVER_PORT}/api/v1/upload_file/example.txt", files={"uploaded_file": f}
+        )
+        assert response.json() == "Successfully uploaded 'example.txt' to the Drive"
+        drive = Drive("lit://uploaded_files", component_name="dummy")
+        assert drive.list() == ["example.txt"]
+        example_path = os.path.join(
+            shared_storage_path(), "artifacts", "drive", "uploaded_files", "file_server", "example.txt"
+        )
+        assert os.path.exists(example_path)
+        with open(example_path) as f:
+            has_found = False
+            lines = f.readlines()
+            assert len(lines) > 0
+            for line in lines:
+                if "test_configure_api_and_upload_files" in line:
+                    has_found = True
+            assert has_found
+
     response = requests.post(
         f"http://localhost:{APP_SERVER_PORT}/api/v1/request", data=InputRequestModel(name="hello").json()
     )