Raven-5 / app.py
WarpWing's picture
Update app.py
c44fcb7 verified
from fastapi import FastAPI, UploadFile, File, HTTPException, Depends
from fastapi.responses import FileResponse
from fastapi.security import OAuth2PasswordBearer
from pathlib import Path
import psutil
import shutil
import os
app = FastAPI()
# Directory where files will be uploaded
UPLOAD_DIRECTORY = Path("uploads")
UPLOAD_DIRECTORY.mkdir(parents=True, exist_ok=True)
# OAuth2 scheme for token-based authorization
oauth2_scheme = OAuth2PasswordBearer(tokenUrl="token")
# Hugging Face secret token from environment variables
VALID_TOKEN = os.environ.get("SECRET_TOKEN")
# Constants to represent the limits for the container resources
MAX_CPU_CORES = 2 # vCPUs
MAX_MEMORY_GB = 16 # GB
MAX_DISK_GB = 50 # GB
# Helper function to calculate the size of a directory
def get_directory_size(directory: Path) -> int:
total_size = 0
for dirpath, dirnames, filenames in os.walk(directory):
for f in filenames:
fp = os.path.join(dirpath, f)
total_size += os.path.getsize(fp)
return total_size
# Dependency to verify the token
def get_current_token(token: str = Depends(oauth2_scheme)):
if token != VALID_TOKEN:
raise HTTPException(status_code=401, detail="Invalid token")
return token
# Health check endpoint
@app.get("/health")
def health_check(token: str = Depends(get_current_token)):
return {
"status": "healthy"
}
# System metrics endpoint: CPU, RAM, and disk usage for uploads folder relative to 2 vCPUs, 16GB RAM, and 50GB disk
@app.get("/metrics")
def get_metrics(token: str = Depends(get_current_token)):
# CPU percentage relative to 2 vCPUs
cpu_percent = round((psutil.cpu_percent(interval=1) / 100) * MAX_CPU_CORES, 2)
# Memory usage relative to 16GB of RAM
memory = psutil.virtual_memory()
memory_used_gb = (memory.total - memory.available) / (1024 ** 3)
# Make sure to cap the used memory at MAX_MEMORY_GB
memory_used_gb_relative = min(memory_used_gb, MAX_MEMORY_GB)
memory_usage_percent = (memory_used_gb_relative / MAX_MEMORY_GB) * 100
# Disk usage for uploads directory relative to 50GB storage
uploads_size_bytes = get_directory_size(UPLOAD_DIRECTORY)
uploads_size_gb = uploads_size_bytes / (1024 ** 3)
uploads_usage_percent = (uploads_size_gb / MAX_DISK_GB) * 100
return {
"cpu_usage_cores": cpu_percent, # Relative CPU usage as cores (out of 2 vCPUs)
"memory": {
"used_gb": round(memory_used_gb_relative, 2), # Capped at 16GB max
"used_percent_of_16gb": round(memory_usage_percent, 2)
},
"disk": {
"uploads_folder_size_gb": round(uploads_size_gb, 2),
"uploads_usage_percent_of_50gb": round(uploads_usage_percent, 2)
}
}
# File upload endpoint
@app.post("/uploadfile/")
async def upload_file(file: UploadFile = File(...), token: str = Depends(get_current_token)):
file_location = UPLOAD_DIRECTORY / file.filename
with file_location.open("wb") as buffer:
shutil.copyfileobj(file.file, buffer)
return {"info": f"file '{file.filename}' saved at '{file_location}'"}
# File download endpoint
@app.get("/downloadfile/{filename}")
def download_file(filename: str, token: str = Depends(get_current_token)):
file_location = UPLOAD_DIRECTORY / filename
if not file_location.exists():
raise HTTPException(status_code=404, detail="File not found")
return FileResponse(path=file_location, filename=filename)
# Show current files endpoint
@app.get("/files/")
def list_files(token: str = Depends(get_current_token)):
files = [f for f in os.listdir(UPLOAD_DIRECTORY) if os.path.isfile(UPLOAD_DIRECTORY / f)]
return {"files": files}