Spaces:
Sleeping
Sleeping
File size: 3,569 Bytes
b1c8f17 eb92a4f 3938059 b5b2e6a b1c8f17 2e76cf7 b1c8f17 c8510e0 75c5f1f |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 |
from fastapi import FastAPI, HTTPException, Request, Query
from fastapi.middleware.cors import CORSMiddleware
from fastapi_cache import FastAPICache
from fastapi_cache.backends.inmemory import InMemoryBackend
from fastapi_cache.decorator import cache
app = FastAPI()
@app.on_event("startup")
async def startup():
FastAPICache.init(InMemoryBackend(), prefix="fastapi-cache")
from pydantic import BaseModel
from typing import List, Dict, Any
TOGETHER_API_KEY = os.getenv('TOGETHER_API_KEY')
BRAVE_API_KEY = os.getenv('BRAVE_API_KEY')
GROQ_API_KEY = os.getenv("GROQ_API_KEY")
HELICON_API_KEY = os.getenv("HELICON_API_KEY")
SUPABASE_USER = os.environ['SUPABASE_USER']
SUPABASE_PASSWORD = os.environ['SUPABASE_PASSWORD']
app.add_middleware(
CORSMiddleware,
allow_origins=["*"],
allow_credentials=True,
allow_methods=["*"],
allow_headers=["*"],)
import asyncio
import uvicorn
from fastapi import FastAPI, HTTPException
from pydantic import BaseModel
from jupyter_client import KernelManager
from typing import Dict
from datetime import datetime, timedelta
import psutil
app = FastAPI()
# Model for input data
class CodeExecutionRequest(BaseModel):
session_token: str
code: str
# Store kernel managers and last access times
kernel_managers: Dict[str, KernelManager] = {}
last_access_times: Dict[str, datetime] = {}
# Timeout duration in seconds
TIMEOUT_DURATION = 600 # 10 minutes
# Function to create a new kernel
async def create_kernel(session_token: str):
km = KernelManager()
km.start_kernel()
kernel_managers[session_token] = km
last_access_times[session_token] = datetime.now()
# Function to kill a kernel
async def kill_kernel(session_token: str):
km = kernel_managers.pop(session_token, None)
if km:
km.shutdown_kernel(now=True)
last_access_times.pop(session_token, None)
# Function to execute code in a kernel
async def execute_code(session_token: str, code: str):
if session_token not in kernel_managers:
await create_kernel(session_token)
km = kernel_managers[session_token]
kc = km.client()
kc.start_channels()
try:
kc.execute(code)
reply = kc.get_shell_msg(timeout=5)
result = reply['content']
except Exception as e:
result = {'status': 'error', 'ename': type(e).__name__, 'evalue': str(e)}
kc.stop_channels()
last_access_times[session_token] = datetime.now()
return result
# Background task to check for idle kernels
async def check_idle_kernels():
while True:
now = datetime.now()
for session_token, last_access in list(last_access_times.items()):
if now - last_access > timedelta(seconds=TIMEOUT_DURATION):
await kill_kernel(session_token)
await asyncio.sleep(60) # Check every minute
@app.on_event("startup")
async def startup_event():
asyncio.create_task(check_idle_kernels())
@app.post("/execute")
async def execute(request: CodeExecutionRequest):
result = await execute_code(request.session_token, request.code)
return result
@app.get("/info")
async def get_info():
# Get the number of active kernels
active_kernels = len(kernel_managers)
# Get system resource usage
cpu_usage = psutil.cpu_percent(interval=1)
ram_usage = psutil.virtual_memory().percent
# Return the information
return {
"active_kernels": active_kernels,
"cpu_usage_percent": cpu_usage,
"ram_usage_percent": ram_usage
}
if __name__ == "__main__":
uvicorn.run(app, host="0.0.0.0", port=8000)
|