moved kernel pre-init to lifespan

Author: mishushakov

template/server/main.py

@@ -34,31 +34,22 @@ async def lifespan(app: FastAPI):
     global client
     client = httpx.AsyncClient()
 
-    with open("/root/.jupyter/kernel_id") as file:
-        default_context_id = file.read().strip()
-
-    default_ws = ContextWebSocket(
-        default_context_id,
-        str(uuid.uuid4()),
-        "python",
-        "/home/user",
-    )
-    default_websockets["python"] = default_context_id
-    websockets["default"] = default_ws
-    websockets[default_context_id] = default_ws
-
-    logger.info("Connecting to default runtime")
-    await default_ws.connect()
-
-    websockets["default"] = default_ws
+    try:
+        default_context = await create_context(client, websockets, "python", "/home/user")
+        default_websockets["python"] = default_context.id
+        websockets["default"] = websockets[default_context.id]
 
-    logger.info("Connected to default runtime")
-    yield
+        logger.info("Connected to default runtime")
+        yield
 
-    for ws in websockets.values():
-        await ws.close()
+        # Will cleanup after application shuts down
+        for ws in websockets.values():
+            await ws.close()
 
-    await client.aclose()
+        await client.aclose()
+    except Exception as e:
+        logger.error(f"Failed to initialize default context: {e}")
+        raise
 
 
 app = FastAPI(lifespan=lifespan)

template/server/messaging.py

@@ -327,15 +327,15 @@ async def _process_message(self, data: dict):
 
         elif data["msg_type"] == "stream":
             if data["content"]["name"] == "stdout":
-                logger.debug(f"Execution {parent_msg_ig} received stdout")
+                logger.debug(f"Execution {parent_msg_ig} received stdout: {data['content']['text']}")
                 await queue.put(
                     Stdout(
                         text=data["content"]["text"], timestamp=data["header"]["date"]
                     )
                 )
 
             elif data["content"]["name"] == "stderr":
-                logger.debug(f"Execution {parent_msg_ig} received stderr")
+                logger.debug(f"Execution {parent_msg_ig} received stderr: {data['content']['text']}")
                 await queue.put(
                     Stderr(
                         text=data["content"]["text"], timestamp=data["header"]["date"]

template/start-up.sh

@@ -13,18 +13,6 @@ function start_jupyter_server() {
 		response=$(curl -s -o /dev/null -w "%{http_code}" "http://localhost:8888/api/status")
 	done
 
-	response=$(curl -s -X POST "localhost:8888/api/sessions" -H "Content-Type: application/json" -d '{"path": "/home/user", "kernel": {"name": "python3"}, "type": "notebook", "name": "default"}')
-	status=$(echo "${response}" | jq -r '.kernel.execution_state')
-	if [[ ${status} != "starting" ]]; then
-		echo "Error creating kernel: ${response} ${status}"
-		exit 1
-	fi
-
-	sudo mkdir -p /root/.jupyter
-	kernel_id=$(echo "${response}" | jq -r '.kernel.id')
-	sudo echo "${kernel_id}" | sudo tee /root/.jupyter/kernel_id >/dev/null
-	sudo echo "${response}" | sudo tee /root/.jupyter/.session_info >/dev/null
-
 	cd /root/.server/
 	/root/.server/.venv/bin/uvicorn main:app --host 0.0.0.0 --port 49999 --workers 1 --no-access-log --no-use-colors --timeout-keep-alive 640
 }