Step 1: Implement the FileServer general structure¶
Let’s dive in on how to develop the component with the following code:
import json
import os
import tarfile
import uuid
import zipfile
from pathlib import Path
from lightning.app import LightningWork, LightningApp
from lightning.app.storage import Drive
class FileServer(LightningWork):
def __init__(
self,
drive: Drive,
base_dir: str = "file_server",
chunk_size=10240,
**kwargs
):
"""This component uploads, downloads files to your application.
Arguments:
drive: The drive can share data inside your application.
base_dir: The local directory where the data will be stored.
chunk_size: The quantity of bytes to download/upload at once.
"""
super().__init__(
cloud_build_config=BuildConfig(["flask, flask-cors"]),
parallel=True,
**kwargs,
)
# 1: Attach the arguments to the state.
self.drive = drive
self.base_dir = base_dir
self.chunk_size = chunk_size
# 2: Create a folder to store the data.
os.makedirs(self.base_dir, exist_ok=True)
# 3: Keep a reference to the uploaded filenames.
def run(self):
# 1: Imports flask requirements.
from flask import Flask, request
from flask_cors import CORS
# 2: Create a flask app
flask_app = Flask(__name__)
CORS(flask_app)
# 3: Define the upload file endpoint
@flask_app.post("/upload_file/")
def upload_file():
"""Upload a file directly as form data."""
f = request.files["file"]
return self.upload_file(f)
@flask_app.get("/")
def list_files():
return self.list_files(str(Path(self.base_dir).resolve()))
# 5: Start the flask app while providing the `host` and `port`.
flask_app.run(host=self.host, port=self.port, load_dotenv=False)
def alive(self):
"""Hack: Returns whether the server is alive."""
return self.url != ""