@@ -0,0 +1,310 @@
# coding=utf-8
"""
@Time :2026/7/27
@Author :tao.chen
"""
import os
import secrets
import socket
import subprocess
import time
import datetime
from contextlib import asynccontextmanager
from pathlib import Path
from typing import Dict , Optional
from fastapi import FastAPI , HTTPException
from pydantic import BaseModel , Field
from loguru import logger
# 全局内存字典：记录运行中的 Jupyter 进程信息
JUPYTER_PROCESSES : Dict [ str , dict ] = { }
WORKSPACES_ROOT = os . getenv ( " WORKSPACES_ROOT " , " /app/workspaces " )
PUBLIC_BASE_URL = os . getenv ( " PUBLIC_BASE_URL " , " http://localhost " )
def get_free_port ( ) - > int :
""" 利用操作系统 socket 特性，动态获取当前闲置的可用端口 """
with socket . socket ( socket . AF_INET , socket . SOCK_STREAM ) as s :
s . bind ( ( " " , 0 ) )
s . listen ( 1 )
port = s . getsockname ( ) [ 1 ]
return port
# 统一请求 Model
class JupyterActionRequest ( BaseModel ) :
action : str = Field (
. . . , description = " 操作类型: ' start ' | ' stop ' | ' list ' "
)
workspace_id : Optional [ str ] = Field (
None , description = " Workspace ID (start 和 stop 操作时必填) "
)
# 辅助处理函数：启动逻辑
def start_process ( cmd , workspace_path , log_dir = " /tmp/process_logs " ) :
log_dir = Path ( log_dir )
log_dir . mkdir ( parents = True , exist_ok = True )
start_time = datetime . datetime . now ( ) . strftime ( " % Y % m %d _ % H % M % S " )
# 临时日志文件
temp_log = log_dir / f " process_start_ { start_time } .log "
log_file = open ( temp_log , " a " , buffering = 1 )
process = subprocess . Popen (
cmd ,
cwd = workspace_path ,
stdout = log_file ,
stderr = subprocess . STDOUT , # stderr 合并到 stdout
start_new_session = True ,
)
# 根据真实 pid 重命名
final_log = log_dir / f " process_ { process . pid } _ { start_time } .log "
log_file . close ( )
temp_log . rename ( final_log )
logger . info (
f " process started pid= { process . pid } , log= { final_log } "
)
return process
def _handle_start ( ws_id : str ) :
workspace_path = Path ( f " ./test/workspaces/ { ws_id } " )
workspace_path . mkdir ( parents = True , exist_ok = True )
# 1. 如果已存在，校验进程状态并复用
if ws_id in JUPYTER_PROCESSES :
p_info = JUPYTER_PROCESSES [ ws_id ]
if p_info [ " process " ] . poll ( ) is None :
logger . info ( f " Workspace { ws_id } already running. " )
return {
" status " : " running " ,
" workspace_id " : ws_id ,
" port " : p_info [ " port " ] ,
" full_url " : p_info [ " full_url " ] ,
}
else :
del JUPYTER_PROCESSES [ ws_id ]
# 2. 动态申请端口与 Token
port = get_free_port ( )
token = secrets . token_hex ( 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=. " ,
# 适用于现代 Jupyter Server / JupyterLab
" --ServerApp.terminals_enabled=False " ,
# 兼容经典 Notebook / 旧版配置项
" --NotebookApp.terminals_enabled=False " ,
]
try :
process = start_process ( cmd , workspace_path . as_posix ( ) )
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 ,
" full_url " : full_url ,
" started_at " : time . time ( ) ,
}
logger . info (
f " Started Jupyter for workspace { ws_id } on port { port } "
)
return {
" pid " : process . pid ,
" base_url " : PUBLIC_BASE_URL ,
" status " : " success " ,
" workspace_id " : ws_id ,
" port " : port ,
" token " : token ,
}
except Exception as e :
logger . error ( f " Failed to start Jupyter for { ws_id } : { str ( e ) } " )
raise HTTPException (
status_code = 500 , detail = f " Failed to start Jupyter: { str ( e ) } "
)
# 辅助处理函数：停止逻辑
def _handle_stop ( ws_id : str ) :
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... "
)
process . kill ( )
process . wait ( )
del JUPYTER_PROCESSES [ ws_id ]
return {
" status " : " stopped " ,
" workspace_id " : ws_id ,
" message " : " Jupyter process terminated and port released. " ,
}
# 辅助处理函数：列表查询逻辑
def _handle_list ( ) :
active_instances = { }
for ws_id , info in list ( JUPYTER_PROCESSES . items ( ) ) :
is_alive = info [ " process " ] . poll ( ) is None
active_instances [ ws_id ] = {
" port " : info [ " port " ] ,
" full_url " : info [ " full_url " ] ,
" is_alive " : is_alive ,
}
return { " status " : " success " , " instances " : active_instances }
def _handle_get ( ws_id : str ) :
""" 【新增】获取指定 Workspace 的 Jupyter 运行状态与完整 URL """
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 :
# 进程如果挂了，清理内存字典并报 404
del JUPYTER_PROCESSES [ ws_id ]
raise HTTPException (
status_code = 404 ,
detail = f " Jupyter process for workspace ' { ws_id } ' has terminated unexpectedly. " ,
)
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 " ] ,
" started_at " : p_info [ " started_at " ] ,
}
# ==================== FastAPI Lifespan 定义 ====================
@asynccontextmanager
async def lifespan ( app : FastAPI ) :
# ---------------- 1. 启动阶段 (Startup) ----------------
logger . info (
f " Initializing workspace processes from root: ' { WORKSPACES_ROOT } ' ... "
)
os . makedirs ( WORKSPACES_ROOT , exist_ok = True )
# 遍历总路径下的每一个子目录（每个子目录名即 workspace_id）
for entry in os . listdir ( WORKSPACES_ROOT ) :
full_path = os . path . join ( WORKSPACES_ROOT , entry )
if os . path . isdir ( full_path ) :
logger . info ( f " Auto-starting Jupyter for workspace ID: ' { entry } ' " )
try :
_handle_start ( entry )
except Exception as err :
logger . error (
f " Startup failed for workspace ' { entry } ' : { err } "
)
yield # 服务在此处保持正常运行，等待 API 请求
# ---------------- 2. 停止阶段 (Shutdown) ----------------
logger . info (
" Service is shutting down. Terminating all active Jupyter sub-processes... "
)
# 复制 key 列表，防止字典在迭代清理过程中因大小改变报错
active_workspaces = list ( JUPYTER_PROCESSES . keys ( ) )
for ws_id in active_workspaces :
try :
_handle_stop ( ws_id )
except Exception as err :
logger . error (
f " Error terminating Jupyter process for ' { ws_id } ' : { err } "
)
logger . info ( " All Jupyter sub-processes have been clean terminated. " )
app = FastAPI ( lifespan = lifespan )
# ---------------- 统一入口 POST 接口 ----------------
@app.get ( " /api/v1/health " )
def healthz ( ) :
return { " status " : " ok " }
@app.post ( " /api/v1/jupyter " )
def handle_jupyter_action ( req : JupyterActionRequest ) :
action = req . action . lower ( )
# 1. 启动操作
if action == " start " :
if not req . workspace_id :
raise HTTPException (
status_code = 400 ,
detail = " ' workspace_id ' is required when action= ' start ' " ,
)
return _handle_start ( req . workspace_id )
# 2. 停止操作
elif action == " stop " :
if not req . workspace_id :
raise HTTPException (
status_code = 400 ,
detail = " ' workspace_id ' is required when action= ' stop ' " ,
)
return _handle_stop ( req . workspace_id )
# 3. 列表操作
elif action == " list " :
return _handle_list ( )
elif action == " get " :
if not req . workspace_id :
raise HTTPException (
status_code = 400 ,
detail = " ' workspace_id ' is required for action= ' get ' " ,
)
return _handle_get ( req . workspace_id )
# 未知操作
else :
raise HTTPException (
status_code = 400 ,
detail = f " Invalid action ' { req . action } ' . Supported actions: ' start ' , ' stop ' , ' list ' " ,
)
tao.chen