Update main.py

main.py

@@ -5,17 +5,24 @@ import shutil
 from fastapi import FastAPI, HTTPException, UploadFile, File
 from fastapi.middleware.cors import CORSMiddleware
 from pydantic import BaseModel
-from jupyter_client.manager import KernelManager  # Updated import
-from typing import Dict
+from jupyter_client.manager import KernelManager
+from typing import Dict, List
 from datetime import datetime, timedelta
 import psutil
-from typing import List
+
+# Set environment variables for Jupyter
+os.environ['JUPYTER_RUNTIME_DIR'] = '/app/kernels'
+os.environ['IPYTHONDIR'] = '/app/ipython'
+os.environ['MPLCONFIGDIR'] = '/app/matplotlib'
 
 app = FastAPI()
 
 # Define root directory for all session folders
 root_dir = os.path.abspath(os.path.dirname(__file__))
 
+# Ensure base output directory exists
+os.makedirs(os.path.join(root_dir, "output"), mode=0o777, exist_ok=True)
+
 # Middleware for CORS
 app.add_middleware(
     CORSMiddleware,
@@ -39,13 +46,14 @@ TIMEOUT_DURATION = 600  # 10 minutes
 
 # Function to create a new kernel and session directory
 async def create_kernel(session_token: str):
-    # Create session directory
     session_dir = os.path.join(root_dir, "output", session_token)
-    if not os.path.exists(session_dir):
-        os.makedirs(session_dir, mode=0o777, exist_ok=True)
+    os.makedirs(session_dir, mode=0o777, exist_ok=True)
     
+    # Create a kernel with specific working directory
     km = KernelManager()
-    km.start_kernel()
+    km.kernel_cmd = ['python', '-c', f'import os; os.chdir("{session_dir}"); from ipykernel.kernelapp import main; main()']
+    km.start_kernel(cwd=session_dir)
+    
     kernel_managers[session_token] = km
     last_access_times[session_token] = datetime.now()
 
@@ -64,34 +72,23 @@ async def kill_kernel(session_token: str):
 @app.post("/upload/{session_token}")
 async def upload_files(session_token: str, files: List[UploadFile] = File(...)):
     session_dir = os.path.join(root_dir, "output", session_token)
-    
-    # Create session directory if it doesn't exist
-    if not os.path.exists(session_dir):
-        os.makedirs(session_dir)
+    os.makedirs(session_dir, mode=0o777, exist_ok=True)
     
     uploaded_files = []
     for file in files:
-        # Save the uploaded file
         file_path = os.path.join(session_dir, file.filename)
         with open(file_path, "wb") as buffer:
             shutil.copyfileobj(file.file, buffer)
         uploaded_files.append(file.filename)
+        os.chmod(file_path, 0o777)  # Give full permissions to uploaded file
 
-    print("uploaded_files",uploaded_files)
+    print("uploaded_files", uploaded_files)
     return {"filenames": uploaded_files, "status": "success"}
 
 # Function to execute code in a kernel
 async def execute_code(session_token: str, code: str):
-    setup_code = """
-%matplotlib inline
-import os
-os.chdir('{session_dir}')
-""".format(session_dir=os.path.join(root_dir, "output", session_token).replace('\\', '\\\\'))
-    
     session_dir = os.path.join(root_dir, "output", session_token)
-    
-    if not os.path.exists(session_dir):
-        os.makedirs(session_dir)
+    os.makedirs(session_dir, mode=0o777, exist_ok=True)
     
     if session_token not in kernel_managers:
         await create_kernel(session_token)
@@ -100,14 +97,18 @@ os.chdir('{session_dir}')
     kc = km.client()
 
     try:
-        # Set working directory for the Python process
-        os.chdir(session_dir)
-        print("Current working directory:", os.getcwd())
+        # Setup code to ensure correct working directory
+        setup_code = f"""
+import os
+os.chdir(r'{session_dir}')
+print('Current working directory:', os.getcwd())
+"""
+        # Execute setup code first
+        kc.execute(setup_code)
+        msg = kc.get_shell_msg()  # Wait for setup completion
         
-        # Execute setup code to set kernel's working directory
-        kc.execute_interactive(setup_code, store_history=False)
         # Execute the provided code
-        kc.execute(code, store_history=False)
+        kc.execute(code)
 
         output = []
         timeout = datetime.now() + timedelta(seconds=TIMEOUT_DURATION)
@@ -153,6 +154,12 @@ async def check_idle_kernels():
 
 @app.on_event("startup")
 async def startup_event():
+    # Ensure all required directories exist with proper permissions
+    os.makedirs('/app/kernels', mode=0o777, exist_ok=True)
+    os.makedirs('/app/ipython', mode=0o777, exist_ok=True)
+    os.makedirs('/app/matplotlib', mode=0o777, exist_ok=True)
+    os.makedirs(os.path.join(root_dir, "output"), mode=0o777, exist_ok=True)
+    
     asyncio.create_task(check_idle_kernels())
 
 @app.post("/execute")
@@ -169,7 +176,6 @@ async def get_info():
     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,
@@ -177,4 +183,4 @@ async def get_info():
     }
 
 if __name__ == "__main__":
-    uvicorn.run(app, host="0.0.0.0", port=7860)
+    uvicorn.run(app, host="0.0.0.0", port=7860)