Update to FastAPI download server

Dockerfile

@@ -1,24 +1,14 @@
-FROM ghcr.io/livebook-dev/livebook:latest-cuda12
+FROM python:3.10-slim
 
-ENV LIVEBOOK_APP_SERVICE_NAME="🐳 Hugging Face - $SPACE_TITLE"
-ENV LIVEBOOK_APP_SERVICE_URL="https://huggingface.co/spaces/$SPACE_AUTHOR_NAME/$SPACE_REPO_NAME"
-ENV LIVEBOOK_UPDATE_INSTRUCTIONS_URL="https://livebook.dev"
-ENV LIVEBOOK_WITHIN_IFRAME="true"
-ENV LIVEBOOK_APPS_PATH="/public-apps"
-ENV LIVEBOOK_APPS_PATH_WARMUP="manual"
-ENV LIVEBOOK_DATA_PATH="/data"
-ENV LIVEBOOK_PORT="7860"
+WORKDIR /app
 
-EXPOSE 7860
+COPY requirements.txt .
+RUN pip install --no-cache-dir -r requirements.txt
 
-RUN mkdir -p /data
-RUN chmod 777 /data
+COPY app.py .
 
-# The Space container runs with user ID 1000, which corresponds to
-# ubuntu in the Livebook image.
-ENV HOME="/home/ubuntu"
+RUN mkdir -p /data && chmod 777 /data
 
-USER ubuntu
+EXPOSE 7860
 
-COPY --chown=ubuntu public-apps/ /public-apps
-RUN /app/bin/warmup_apps
+CMD ["uvicorn", "app:app", "--host", "0.0.0.0", "--port", "7860"]

app.py

@@ -0,0 +1,129 @@
+import asyncio
+import os
+import time
+import uuid
+import aiohttp
+from fastapi import FastAPI, BackgroundTasks, Request
+from fastapi.responses import FileResponse, StreamingResponse
+from pydantic import BaseModel
+
+app = FastAPI()
+
+# In-memory store for download tasks
+# tasks = { "task_id": { "url": str, "status": "downloading"|"completed"|"error", "total_size": int, "downloaded": int, "speed": float, "file_path": str } }
+tasks = {}
+
+DATA_DIR = "/data"
+os.makedirs(DATA_DIR, exist_ok=True)
+
+class DownloadRequest(BaseModel):
+    url: str
+
+async def download_file(task_id: str, url: str):
+    file_path = os.path.join(DATA_DIR, f"{task_id}.bin")
+    tasks[task_id] = {
+        "url": url,
+        "status": "downloading",
+        "total_size": 0,
+        "downloaded": 0,
+        "speed": 0.0,
+        "file_path": file_path
+    }
+    
+    try:
+        async with aiohttp.ClientSession() as session:
+            async with session.get(url) as response:
+                response.raise_for_status()
+                total_size = int(response.headers.get('Content-Length', 0))
+                tasks[task_id]["total_size"] = total_size
+                
+                downloaded = 0
+                start_time = time.time()
+                last_time = start_time
+                last_downloaded = 0
+                
+                with open(file_path, 'wb') as f:
+                    async for chunk in response.content.iter_chunked(1024 * 1024): # 1MB chunks to save memory
+                        if not chunk:
+                            break
+                        f.write(chunk)
+                        downloaded += len(chunk)
+                        tasks[task_id]["downloaded"] = downloaded
+                        
+                        current_time = time.time()
+                        if current_time - last_time >= 1.0: # Update speed every second
+                            speed = (downloaded - last_downloaded) / (current_time - last_time)
+                            tasks[task_id]["speed"] = speed
+                            last_time = current_time
+                            last_downloaded = downloaded
+                            
+                tasks[task_id]["status"] = "completed"
+                tasks[task_id]["speed"] = 0.0
+    except Exception as e:
+        tasks[task_id]["status"] = "error"
+        tasks[task_id]["error"] = str(e)
+
+@app.post("/start_download")
+async def start_download(req: DownloadRequest, background_tasks: BackgroundTasks):
+    task_id = str(uuid.uuid4())
+    background_tasks.add_task(download_file, task_id, req.url)
+    return {"task_id": task_id}
+
+@app.get("/status/{task_id}")
+async def get_status(task_id: str):
+    if task_id not in tasks:
+        return {"error": "Task not found"}
+    return tasks[task_id]
+
+@app.get("/download/{task_id}")
+async def download(task_id: str):
+    if task_id not in tasks or tasks[task_id]["status"] != "completed":
+        return {"error": "File not ready"}
+    return FileResponse(tasks[task_id]["file_path"], filename=f"downloaded_{task_id}.bin")
+
+@app.get("/stream/{task_id}")
+async def stream(task_id: str, request: Request):
+    if task_id not in tasks or tasks[task_id]["status"] != "completed":
+        return {"error": "File not ready"}
+    
+    file_path = tasks[task_id]["file_path"]
+    file_size = os.path.getsize(file_path)
+    
+    range_header = request.headers.get("Range")
+    if range_header:
+        # Simple range request handling for streaming
+        byte1, byte2 = 0, None
+        match = range_header.replace("bytes=", "").split("-")
+        if match[0]:
+            byte1 = int(match[0])
+        if len(match) > 1 and match[1]:
+            byte2 = int(match[1])
+        
+        length = file_size - byte1
+        if byte2 is not None:
+            length = byte2 + 1 - byte1
+            
+        def file_iterator(start, length):
+            with open(file_path, "rb") as f:
+                f.seek(start)
+                chunk_size = 1024 * 1024
+                while length > 0:
+                    read_size = min(chunk_size, length)
+                    data = f.read(read_size)
+                    if not data:
+                        break
+                    yield data
+                    length -= len(data)
+                    
+        headers = {
+            "Content-Range": f"bytes {byte1}-{byte1+length-1}/{file_size}",
+            "Accept-Ranges": "bytes",
+            "Content-Length": str(length),
+        }
+        return StreamingResponse(file_iterator(byte1, length), status_code=206, headers=headers)
+    else:
+        return FileResponse(file_path)
+
+if __name__ == "__main__":
+    import uvicorn
+    uvicorn.run(app, host="0.0.0.0", port=7860)

requirements.txt

@@ -0,0 +1,4 @@
+fastapi
+uvicorn
+aiohttp
+pydantic