269 lines
8.4 KiB
Python
269 lines
8.4 KiB
Python
"""Jupyter process registry and lifecycle.
|
|
|
|
Owns the in-memory ``JUPYTER_PROCESSES`` dict, the ``STATE_LOCK`` that
|
|
serializes mutations to it, and the per-workspace start/stop/list/get
|
|
operations. Workspace discovery (startup scan) lives here because it is
|
|
a thin wrapper over ``start_workspace``.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import asyncio
|
|
import os
|
|
import secrets
|
|
import subprocess
|
|
import time
|
|
from pathlib import Path
|
|
from typing import TypedDict
|
|
|
|
from fastapi import HTTPException
|
|
from loguru import logger
|
|
|
|
from common.config import settings
|
|
from common.utils import get_free_port, start_process
|
|
from runtime.mount import WORKSPACES_ROOT
|
|
|
|
PUBLIC_BASE_URL = settings.public_base_url
|
|
|
|
|
|
class JupyterProcessRecord(TypedDict):
|
|
process: subprocess.Popen
|
|
port: int
|
|
token: str
|
|
base_url: str
|
|
started_at: float
|
|
|
|
|
|
JUPYTER_PROCESSES: dict[str, JupyterProcessRecord] = {}
|
|
WORKSPACE_LOCKS: dict[str, asyncio.Lock] = {}
|
|
_LOCKS_REGISTRY = asyncio.Lock()
|
|
|
|
|
|
def get_workspace_lock(ws_id: str) -> asyncio.Lock:
|
|
"""Return the per-workspace lock, creating it on first use.
|
|
|
|
``asyncio.Lock`` is created lazily per event loop; sharing it across
|
|
loops is unsafe. The instance lives for the lifetime of the process.
|
|
"""
|
|
if ws_id in WORKSPACE_LOCKS:
|
|
return WORKSPACE_LOCKS[ws_id]
|
|
# We do not need to hold the registry lock long enough to block; a
|
|
# double-check pattern prevents accidentally replacing an existing
|
|
# lock. The Lock object is itself safe to call .locked()/acquire on.
|
|
lock = WORKSPACE_LOCKS.get(ws_id)
|
|
if lock is None:
|
|
lock = asyncio.Lock()
|
|
WORKSPACE_LOCKS[ws_id] = lock
|
|
return lock
|
|
|
|
|
|
def _drop_workspace_lock(ws_id: str) -> None:
|
|
"""Remove the per-workspace lock once no process references it.
|
|
|
|
Called after a successful stop to keep the registry bounded. We only
|
|
drop a lock we created and only when it is not held.
|
|
"""
|
|
lock = WORKSPACE_LOCKS.get(ws_id)
|
|
if lock is None or lock.locked():
|
|
return
|
|
WORKSPACE_LOCKS.pop(ws_id, None)
|
|
|
|
|
|
async def start_workspace(ws_id: str) -> dict:
|
|
async with get_workspace_lock(ws_id):
|
|
workspace_path = WORKSPACES_ROOT / ws_id
|
|
|
|
if ws_id in JUPYTER_PROCESSES:
|
|
p_info = JUPYTER_PROCESSES[ws_id]
|
|
if p_info["process"].poll() is None:
|
|
if time.time() - p_info["started_at"] > 24 * 3600:
|
|
logger.warning(
|
|
f"Reusing Jupyter for {ws_id} older than 24h "
|
|
f"(started_at={p_info['started_at']})"
|
|
)
|
|
return {
|
|
"status": "running",
|
|
"workspace_id": ws_id,
|
|
"pid": p_info["process"].pid,
|
|
"port": p_info["port"],
|
|
"token": p_info["token"],
|
|
"base_url": PUBLIC_BASE_URL,
|
|
"full_url": (
|
|
f"{PUBLIC_BASE_URL}:{p_info['port']}/jupyter/{ws_id}/"
|
|
f"?token={p_info['token']}"
|
|
),
|
|
}
|
|
del JUPYTER_PROCESSES[ws_id]
|
|
|
|
port = get_free_port()
|
|
token = secrets.token_urlsafe(16)
|
|
base_path = f"/jupyter/{ws_id}/"
|
|
|
|
cmd = [
|
|
"jupyter", "notebook",
|
|
f"--port={port}",
|
|
"--ip=0.0.0.0",
|
|
"--no-browser",
|
|
"--allow-root",
|
|
f"--ServerApp.token={token}",
|
|
f"--ServerApp.base_url={base_path}",
|
|
"--notebook-dir=.",
|
|
"--ServerApp.terminals_enabled=False",
|
|
"--NotebookApp.terminals_enabled=False",
|
|
"--ServerApp.allow_origin=*",
|
|
"--NotebookApp.allow_origin=*",
|
|
"--ServerApp.disable_check_xsrf=True",
|
|
"--NotebookApp.disable_check_xsrf=True",
|
|
]
|
|
|
|
try:
|
|
process, log_file = start_process(cmd, workspace_path)
|
|
except Exception as e:
|
|
logger.error(f"Failed to start Jupyter for {ws_id}: {e}")
|
|
raise HTTPException(
|
|
status_code=500, detail=f"Failed to start Jupyter: {e}"
|
|
)
|
|
|
|
full_url = f"{PUBLIC_BASE_URL}:{port}{base_path}?token={token}"
|
|
|
|
JUPYTER_PROCESSES[ws_id] = {
|
|
"process": process,
|
|
"base_url": PUBLIC_BASE_URL,
|
|
"port": port,
|
|
"token": token,
|
|
"started_at": time.time(),
|
|
}
|
|
|
|
logger.info(
|
|
f"Started Jupyter for workspace {ws_id} "
|
|
f"pid={process.pid} port={port} log={log_file}"
|
|
)
|
|
return {
|
|
"pid": process.pid,
|
|
"base_url": PUBLIC_BASE_URL,
|
|
"status": "success",
|
|
"workspace_id": ws_id,
|
|
"port": port,
|
|
"token": token,
|
|
"full_url": full_url,
|
|
}
|
|
|
|
|
|
async def stop_workspace(ws_id: str) -> dict:
|
|
async with get_workspace_lock(ws_id):
|
|
if ws_id not in JUPYTER_PROCESSES:
|
|
raise HTTPException(
|
|
status_code=404,
|
|
detail=f"No active Jupyter process found for workspace '{ws_id}'",
|
|
)
|
|
|
|
p_info = JUPYTER_PROCESSES[ws_id]
|
|
process: subprocess.Popen = p_info["process"]
|
|
|
|
if process.poll() is None:
|
|
try:
|
|
process.terminate()
|
|
process.wait(timeout=3)
|
|
logger.info(f"Gracefully stopped Jupyter for workspace {ws_id}")
|
|
except subprocess.TimeoutExpired:
|
|
logger.warning(
|
|
f"Jupyter for {ws_id} did not stop gracefully. Force killing..."
|
|
)
|
|
try:
|
|
process.kill()
|
|
process.wait()
|
|
except Exception as err:
|
|
logger.error(f"Failed to kill Jupyter process for {ws_id}: {err}")
|
|
|
|
del JUPYTER_PROCESSES[ws_id]
|
|
|
|
_drop_workspace_lock(ws_id)
|
|
return {
|
|
"status": "stopped",
|
|
"workspace_id": ws_id,
|
|
"message": "Jupyter process terminated and port released.",
|
|
}
|
|
|
|
|
|
async def list_workspaces() -> dict:
|
|
async with _LOCKS_REGISTRY:
|
|
snapshot = dict(JUPYTER_PROCESSES)
|
|
return {
|
|
"status": "success",
|
|
"instances": {
|
|
ws_id: {
|
|
"port": info["port"],
|
|
"full_url": (
|
|
f"{PUBLIC_BASE_URL}:{info['port']}/jupyter/{ws_id}/"
|
|
f"?token={info['token']}"
|
|
),
|
|
"is_alive": info["process"].poll() is None,
|
|
}
|
|
for ws_id, info in snapshot.items()
|
|
},
|
|
}
|
|
|
|
|
|
async def get_workspace(ws_id: str) -> dict:
|
|
async with get_workspace_lock(ws_id):
|
|
if ws_id not in JUPYTER_PROCESSES:
|
|
raise HTTPException(
|
|
status_code=404,
|
|
detail=f"No active Jupyter process found for workspace '{ws_id}'",
|
|
)
|
|
|
|
p_info = JUPYTER_PROCESSES[ws_id]
|
|
is_alive = p_info["process"].poll() is None
|
|
|
|
if not is_alive:
|
|
del JUPYTER_PROCESSES[ws_id]
|
|
else:
|
|
return {
|
|
"status": "running",
|
|
"pid": p_info["process"].pid,
|
|
"base_url": PUBLIC_BASE_URL,
|
|
"workspace_id": ws_id,
|
|
"port": p_info["port"],
|
|
"token": p_info["token"],
|
|
"full_url": (
|
|
f"{PUBLIC_BASE_URL}:{p_info['port']}/jupyter/{ws_id}/"
|
|
f"?token={p_info['token']}"
|
|
),
|
|
"started_at": p_info["started_at"],
|
|
}
|
|
|
|
_drop_workspace_lock(ws_id)
|
|
raise HTTPException(
|
|
status_code=404,
|
|
detail=(
|
|
f"Jupyter process for workspace '{ws_id}' "
|
|
"has terminated unexpectedly."
|
|
),
|
|
)
|
|
|
|
|
|
async def scan_workspaces() -> None:
|
|
if not WORKSPACES_ROOT.exists():
|
|
return
|
|
try:
|
|
entries = os.listdir(WORKSPACES_ROOT)
|
|
except Exception as e:
|
|
logger.error(f"scan workspace failed: {e}")
|
|
return
|
|
|
|
sem = asyncio.Semaphore(4)
|
|
|
|
async def _start(entry: str) -> None:
|
|
path = WORKSPACES_ROOT / entry
|
|
if not path.is_dir():
|
|
return
|
|
logger.info(f"Found workspace: {entry}")
|
|
logger.info(f"Auto-starting Jupyter for workspace ID: '{entry}'")
|
|
async with sem:
|
|
try:
|
|
await start_workspace(entry)
|
|
except Exception as err:
|
|
logger.error(f"Startup failed for workspace '{entry}': {err}")
|
|
|
|
await asyncio.gather(*[_start(entry) for entry in entries])
|