Spaces:
Sleeping
Sleeping
File size: 4,504 Bytes
b1c8f17 eb92a4f 9ff2bfb 5dbbbb0 3938059 9ff2bfb 3938059 b5b2e6a b1c8f17 2e76cf7 b1c8f17 c8510e0 75c5f1f ee3a211 75c5f1f ee3a211 e82217e ee3a211 75c5f1f ee3a211 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 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 |
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
import os
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):
setup_code = "%matplotlib inline"
if session_token not in kernel_managers:
await create_kernel(session_token)
km = kernel_managers[session_token]
kc = km.client()
kc.start_channels()
try:
# Execute setup code
kc.execute_interactive(setup_code, store_history=False)
# Execute the provided code
kc.execute_interactive(code, store_history=False)
output = []
while True:
msg = kc.get_iopub_msg()
if msg['msg_type'] == 'status' and msg['content']['execution_state'] == 'idle':
break
elif msg['msg_type'] == 'error':
error_output = {
"ename": msg['content']['ename'],
"evalue": msg['content']['evalue'],
"traceback": msg['content']['traceback']
}
output.append({"error": error_output})
if 'data' in msg['content']:
output.append({"data": msg['content']['data']})
elif 'text' in msg['content']:
output.append({"text": msg['content']['text']})
return {'status': 'success', 'value': output}
except Exception as e:
result = {'status': 'error', 'value': 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)
|