diff --git a/src/lightning_app/CHANGELOG.md b/src/lightning_app/CHANGELOG.md index 45cfa83166ea4..e7ca605b4c2b5 100644 --- a/src/lightning_app/CHANGELOG.md +++ b/src/lightning_app/CHANGELOG.md @@ -15,6 +15,9 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/). - Add `--secret` option to CLI to allow binding Secrets to app environment variables when running in the cloud ([#14612](https://github.com/Lightning-AI/lightning/pull/14612)) +- 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)) diff --git a/src/lightning_app/core/api.py b/src/lightning_app/core/api.py index a1b75dcea1383..7ba6b3de104cc 100644 --- a/src/lightning_app/core/api.py +++ b/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: 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 diff --git a/src/lightning_app/runners/cloud.py b/src/lightning_app/runners/cloud.py index b873f6db48ab6..fe7c9e4b8751a 100644 --- a/src/lightning_app/runners/cloud.py +++ b/src/lightning_app/runners/cloud.py @@ -351,7 +351,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