848 lines
30 KiB
Python
848 lines
30 KiB
Python
"""
|
||
app/api/modeling.py
|
||
===================
|
||
|
||
机器学习与水质反演相关的 API 路由。
|
||
|
||
接口（最终路径, 挂载后）：
|
||
POST /api/modeling/train 提交模型训练任务, 立即返回 task_id
|
||
GET /api/modeling/models 列出已训练好的模型（未来从磁盘 joblib 读）
|
||
POST /api/modeling/predict 提交模型推断任务, 立即返回 task_id
|
||
|
||
设计要点
|
||
--------
|
||
- 训练 / 推断均为异步后台任务, 复用 app.core.task_store 的并发安全任务状态。
|
||
- 模型元数据用模块级 _MODEL_REGISTRY 暂存（开发期内存存储）,
|
||
未来从磁盘 joblib 读时只需替换 list_trained_models() 内部实现即可。
|
||
- /predict 已接入真实 sklearn + xarray + dask 流式推断:
|
||
* joblib.load 读模型（缺文件时降级为 Dummy RandomForestRegressor）
|
||
* xr.open_zarr 延迟打开影像, NaN 填 0
|
||
* xr.apply_ufunc(dask="parallelized") 沿 (y, x) 逐 chunk 调 model.predict
|
||
* to_zarr(mode="w", compute=True) 流式写出, 内存峰值 ≈ 1 个 chunk
|
||
- /train 已接入真实 sklearn + pandas 训练管线:
|
||
* pd.read_csv 读结构化训练表（支持 [lat, lon, target_*, feature_*] 布局）
|
||
* 按 target 列 dropna 清洗；按 feature_start 索引/列名切分特征
|
||
* sklearn Pipeline: StandardScaler -> {RF/SVR/LinearRegression/KNN/PLS}
|
||
* train_test_split(80/20) 划分, 计算 test_r2/rmse/mae
|
||
* joblib.dump({model, metadata}) 落盘 ./data/models/{model_id}.joblib
|
||
* 测试指标写回 TASK_STORE, 同时登记到 _MODEL_REGISTRY
|
||
注: 旧版 SPXY / KS 划分留作未来扩展, 当前固定 random 划分 (test_size=0.2, random_state=42)。
|
||
"""
|
||
|
||
import asyncio
|
||
import shutil
|
||
import traceback
|
||
import uuid
|
||
from datetime import datetime
|
||
from pathlib import Path
|
||
from typing import Any, Dict, List, Optional, Union
|
||
|
||
import joblib
|
||
import numpy as np
|
||
import pandas as pd
|
||
import xarray as xr
|
||
from fastapi import APIRouter, BackgroundTasks, HTTPException, UploadFile, File
|
||
from pydantic import BaseModel, Field
|
||
from sklearn.cross_decomposition import PLSRegression
|
||
from sklearn.ensemble import RandomForestRegressor
|
||
from sklearn.linear_model import LinearRegression
|
||
from sklearn.metrics import mean_absolute_error, mean_squared_error, r2_score
|
||
from sklearn.model_selection import train_test_split
|
||
from sklearn.neighbors import KNeighborsRegressor
|
||
from sklearn.pipeline import Pipeline
|
||
from sklearn.preprocessing import StandardScaler
|
||
from sklearn.svm import SVR
|
||
|
||
# 复用并发安全任务状态存储（与 deglint 共享同一份 TASK_STORE,
|
||
# 通过 task 记录里的 "kind" 字段区分 train / predict / deglint）
|
||
from app.core.task_store import get_task, set_task, update_task
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# 路由实例
|
||
# ---------------------------------------------------------------------------
|
||
# prefix="/modeling" 让本文件内只写 /train /models /predict 等短路径,
|
||
# 最终完整路径由 main.py 挂载时再补 /api。
|
||
# ---------------------------------------------------------------------------
|
||
router = APIRouter(prefix="/modeling", tags=["modeling"])
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# 数据模型
|
||
# ---------------------------------------------------------------------------
|
||
class TrainRequest(BaseModel):
|
||
"""POST /api/modeling/train 的请求体"""
|
||
|
||
model_type: str = Field(
|
||
...,
|
||
description="模型类型, 例如 'RF' (随机森林) / 'SVR' (支持向量回归) / 'XGBoost' / 'MLP'",
|
||
examples=["RF", "SVR"],
|
||
)
|
||
target: str = Field(
|
||
...,
|
||
description="反演目标水质参数, 例如 'Chl-a' (叶绿素a) / 'TSS' (总悬浮物) / 'CDOM' (有色可溶有机物)",
|
||
examples=["Chl-a", "TSS", "CDOM"],
|
||
)
|
||
train_data_path: str = Field(
|
||
...,
|
||
description="训练数据集的 zarr 路径（包含 reflectance 变量与 target 标签）",
|
||
examples=["./data/train.zarr"],
|
||
)
|
||
feature_start: Union[int, str] = Field(
|
||
default=4,
|
||
description=(
|
||
"特征列起始位置. 表格布局假定为 "
|
||
"[lat, lon, target_1, target_2, ..., feature_1, feature_2, ...] "
|
||
"可传 int 列索引（如 4）或 str 列名（如 '374.285' 波长起点）。"
|
||
"默认 4, 即前 4 列视为元数据/目标, 之后全部是特征。"
|
||
),
|
||
examples=[4, "374.285"],
|
||
)
|
||
params: Dict[str, Any] = Field(
|
||
default_factory=dict,
|
||
description="模型超参, 例如 RF 的 {'n_estimators': 100, 'max_depth': 20}",
|
||
examples=[{"n_estimators": 100, "max_depth": 20}],
|
||
)
|
||
|
||
|
||
class PredictRequest(BaseModel):
|
||
"""POST /api/modeling/predict 的请求体"""
|
||
|
||
model_id: str = Field(
|
||
...,
|
||
description="已训练模型的 ID（由 /api/modeling/train 返回或 /api/modeling/models 列出）",
|
||
)
|
||
input_zarr_path: str = Field(
|
||
...,
|
||
description="待推断影像的 zarr 路径",
|
||
examples=["./data/scene.zarr"],
|
||
)
|
||
output_zarr_path: Optional[str] = Field(
|
||
default=None,
|
||
description=(
|
||
"输出 zarr 路径, 缺省时由后端按规则生成 "
|
||
"(如 ./data/{model_id}_{input_stem}_pred.zarr)"
|
||
),
|
||
)
|
||
|
||
|
||
class TaskAcceptedResponse(BaseModel):
|
||
"""提交训练/推断任务后立即返回的响应"""
|
||
|
||
task_id: str
|
||
status: str # 一定是 PENDING
|
||
kind: str # "train" / "predict", 便于前端识别任务类型
|
||
|
||
|
||
class ModelInfo(BaseModel):
|
||
"""单个模型的元信息（GET /api/modeling/models 的元素）"""
|
||
|
||
model_id: str
|
||
model_type: str
|
||
target: str
|
||
params: Dict[str, Any]
|
||
path: str # joblib 文件路径
|
||
created_at: str
|
||
train_task_id: str # 产生此模型的那个训练任务的 ID
|
||
|
||
|
||
class ModelListResponse(BaseModel):
|
||
"""GET /api/modeling/models 的响应"""
|
||
|
||
models: List[ModelInfo]
|
||
count: int
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# 模块级模型注册表（开发期内存, 未来替换为磁盘扫描）
|
||
# ---------------------------------------------------------------------------
|
||
# model_id -> ModelInfo 字典
|
||
# 读多写少, 用一个普通 dict 足够（CPython GIL 兜底）。
|
||
# 写时（训练完成时）只发生一次, 无并发风险。
|
||
# ---------------------------------------------------------------------------
|
||
_MODEL_REGISTRY: Dict[str, Dict[str, Any]] = {}
|
||
|
||
|
||
def _register_model(record: Dict[str, Any]) -> None:
|
||
"""将训练完成的模型登记到内存注册表。"""
|
||
_MODEL_REGISTRY[record["model_id"]] = record
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# 训练管线的模块级辅助函数
|
||
# ---------------------------------------------------------------------------
|
||
# 设计要点 (与推断管线一致):
|
||
# 1) 模块级函数: dask / joblib 后端若走子进程 pickle, 嵌套闭包会丢字段。
|
||
# 2) 同步执行: execute_train_task 用 asyncio.to_thread 派发, 内部全程同步阻塞。
|
||
# 3) 失败抛异常: 异常由 execute_train_task 捕获, 转 FAILED + traceback。
|
||
# ---------------------------------------------------------------------------
|
||
|
||
# model_type (大写字符串) -> sklearn 估计器类
|
||
# 与 OpenClaw model_configs 思路一致, 但此处只保留类 (参数由 params 透传)
|
||
_MODEL_CLASS_REGISTRY: Dict[str, type] = {
|
||
"RF": RandomForestRegressor,
|
||
"SVR": SVR,
|
||
"LinearRegression": LinearRegression,
|
||
"KNN": KNeighborsRegressor,
|
||
"PLS": PLSRegression,
|
||
}
|
||
|
||
|
||
def _get_model_pipeline(model_type: str, params: Optional[Dict[str, Any]]) -> Pipeline:
|
||
"""
|
||
模型工厂: 按 model_type 选 sklearn 类, 用 StandardScaler + 估计器构造 Pipeline。
|
||
|
||
与 OpenClaw 不同之处: 把 scaler 放进 Pipeline 第一步,
|
||
推断时直接 pipeline.predict(X) 即可, scaler 参数与训练时严格一致。
|
||
"""
|
||
model_cls = _MODEL_CLASS_REGISTRY.get(model_type)
|
||
if model_cls is None:
|
||
raise ValueError(
|
||
f"不支持的 model_type='{model_type}', "
|
||
f"可选: {sorted(_MODEL_CLASS_REGISTRY.keys())}"
|
||
)
|
||
estimator = model_cls(**(params or {}))
|
||
return Pipeline([("scaler", StandardScaler()), ("model", estimator)])
|
||
|
||
|
||
def _load_train_df(csv_path: str) -> pd.DataFrame:
|
||
"""
|
||
读 CSV 训练表, 规整空串 / 空白 / NULL 等为 NaN。
|
||
|
||
沿用 OpenClaw modeling_batch.load_data_batch 的读取策略:
|
||
na_values 显式列举 + 正则二次清理 (防 cell 内出现 " " 等纯空白)。
|
||
"""
|
||
try:
|
||
df = pd.read_csv(
|
||
csv_path,
|
||
na_values=["", " ", "NaN", "nan", "NULL", "null"],
|
||
)
|
||
except FileNotFoundError as exc:
|
||
raise FileNotFoundError(f"训练数据文件不存在: {csv_path}") from exc
|
||
except pd.errors.EmptyDataError as exc:
|
||
raise ValueError(f"训练数据文件为空: {csv_path}") from exc
|
||
# 二次清理: 残留的纯空白 cell
|
||
df = df.replace(r"^\s*$", np.nan, regex=True)
|
||
return df
|
||
|
||
|
||
def _resolve_feature_start(
|
||
df: pd.DataFrame,
|
||
feature_start: Union[int, str],
|
||
) -> int:
|
||
"""
|
||
将 feature_start (int 索引 / str 列名) 统一解析为 int 列索引。
|
||
|
||
与 OpenClaw modeling_batch.load_data_batch / load_data_single 一致:
|
||
str 走 columns.get_loc, int 直接返回。
|
||
"""
|
||
if isinstance(feature_start, str):
|
||
if feature_start not in df.columns:
|
||
raise ValueError(
|
||
f"feature_start='{feature_start}' 不在 CSV 列中: {list(df.columns)}"
|
||
)
|
||
return int(df.columns.get_loc(feature_start))
|
||
return int(feature_start)
|
||
|
||
|
||
def _run_train_sync(
|
||
model_type: str,
|
||
target: str,
|
||
train_data_path: str,
|
||
feature_start: Union[int, str],
|
||
params: Optional[Dict[str, Any]],
|
||
output_model_path: Path,
|
||
) -> Dict[str, Any]:
|
||
"""
|
||
完整同步训练流程 (由 execute_train_task 在线程池内调用):
|
||
|
||
pd.read_csv -> 目标列 dropna -> 切特征 -> train_test_split(80/20)
|
||
-> Pipeline(StandardScaler + model).fit -> 评估 test_r2/rmse/mae
|
||
-> joblib.dump({model, metadata}, output_model_path)
|
||
|
||
Returns:
|
||
metadata 字典, 含 test_r2 / test_rmse / test_mae / n_features 等,
|
||
调用方负责写回 TASK_STORE 和 _MODEL_REGISTRY。
|
||
|
||
注: 旧版 SPXY / KS 划分留作未来扩展 (params.split_method 控制),
|
||
当前固定 random + test_size=0.2 + random_state=42。
|
||
"""
|
||
df = _load_train_df(train_data_path)
|
||
|
||
if target not in df.columns:
|
||
raise ValueError(
|
||
f"target='{target}' 不在 CSV 列中, 可选: {list(df.columns)}"
|
||
)
|
||
|
||
# 1) 清洗: 仅剔除 target NaN 的行 (与 OpenClaw load_data_single 一致)
|
||
df = df[df[target].notna()].copy()
|
||
if df.empty:
|
||
raise ValueError("target 剔除 NaN 后无样本, 终止训练")
|
||
|
||
# 2) 特征切分
|
||
feature_start_idx = _resolve_feature_start(df, feature_start)
|
||
feature_columns = list(df.columns[feature_start_idx:])
|
||
|
||
X = df.iloc[:, feature_start_idx:].astype(np.float64)
|
||
y = df[target].astype(np.float64).values
|
||
|
||
# 3) 划分 (固定 random, 未来扩展 spxy/ks)
|
||
X_train, X_test, y_train, y_test = train_test_split(
|
||
X.values,
|
||
y,
|
||
test_size=0.2,
|
||
random_state=42,
|
||
)
|
||
|
||
# 4) 构造 Pipeline + 训练
|
||
pipeline = _get_model_pipeline(model_type, params)
|
||
pipeline.fit(X_train, y_train)
|
||
|
||
# 5) 测试集与训练集评估
|
||
y_pred = pipeline.predict(X_test)
|
||
test_r2 = float(r2_score(y_test, y_pred))
|
||
test_rmse = float(np.sqrt(mean_squared_error(y_test, y_pred)))
|
||
test_mae = float(mean_absolute_error(y_test, y_pred))
|
||
|
||
y_train_pred = pipeline.predict(X_train)
|
||
train_r2 = float(r2_score(y_train, y_train_pred))
|
||
train_rmse = float(np.sqrt(mean_squared_error(y_train, y_train_pred)))
|
||
train_mae = float(mean_absolute_error(y_train, y_train_pred))
|
||
|
||
metadata: Dict[str, Any] = {
|
||
"model_type": model_type,
|
||
"target": target,
|
||
"feature_start": feature_start,
|
||
"feature_columns": feature_columns,
|
||
"n_features": int(X.shape[1]),
|
||
"n_samples": int(X.shape[0]),
|
||
"train_size": int(X_train.shape[0]),
|
||
"test_size": int(X_test.shape[0]),
|
||
"params": dict(params or {}),
|
||
"test_r2": test_r2,
|
||
"test_rmse": test_rmse,
|
||
"test_mae": test_mae,
|
||
"train_r2": train_r2,
|
||
"train_rmse": train_rmse,
|
||
"train_mae": train_mae,
|
||
"split_method": "random",
|
||
"trained_at": datetime.now().isoformat(),
|
||
}
|
||
|
||
# 7) 持久化 (目录可能不存在, 顺手建)
|
||
output_model_path = Path(output_model_path)
|
||
output_model_path.parent.mkdir(parents=True, exist_ok=True)
|
||
joblib.dump(
|
||
{"model": pipeline, "metadata": metadata},
|
||
output_model_path,
|
||
)
|
||
|
||
return metadata
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# 推断管线的模块级辅助函数
|
||
# ---------------------------------------------------------------------------
|
||
# 设计要点:
|
||
# 1) Dask 调度时, 函数必须可被工作进程 pickle 序列化。
|
||
# 因此 _predict_block / _load_model / _make_dummy_model / _run_predict_sync
|
||
# 全部是模块级函数 (而非嵌套), 避免闭包陷阱。
|
||
# 2) _predict_block 通过 model.predict(spectra_2d) 整批预测,
|
||
# 整张影像的 O(n_pixels * n_bands) 一次性预测在大矩阵上必 OOM,
|
||
# 因此外层用 xr.apply_ufunc(dask="parallelized") 把矩阵切块
|
||
# 逐块进入此函数, 单次内存峰值 ≈ 1 个 (y_chunk, x_chunk, band) 大小。
|
||
# ---------------------------------------------------------------------------
|
||
|
||
|
||
def _make_dummy_model(n_features: int) -> RandomForestRegressor:
|
||
"""
|
||
构造一个 Dummy 随机森林回归器。
|
||
|
||
用途:
|
||
1) 真实 joblib 文件不存在时的连通性测试
|
||
2) 训练骨架尚未接入真实数据时的占位推断
|
||
"""
|
||
rng = np.random.default_rng(42)
|
||
X = rng.random((200, n_features))
|
||
y = rng.random(200)
|
||
model = RandomForestRegressor(
|
||
n_estimators=10, max_depth=5, random_state=0, n_jobs=1
|
||
)
|
||
model.fit(X, y)
|
||
return model
|
||
|
||
|
||
def _load_model(path: str, n_features: int) -> Any:
|
||
"""
|
||
加载训练好的 sklearn 模型, 失败时降级 Dummy。
|
||
|
||
优先级:
|
||
1) path 存在且 joblib.load 成功 -> 返回真实模型
|
||
2) 否则 -> 降级为 Dummy 随机森林 (n_features 必须指定)
|
||
"""
|
||
p = Path(path)
|
||
if p.is_file() and p.stat().st_size > 0:
|
||
try:
|
||
print(f"[model] 从磁盘加载: {path}")
|
||
return joblib.load(path)
|
||
except Exception as exc: # noqa: BLE001
|
||
print(f"[model] joblib.load 失败 ({type(exc).__name__}: {exc}), 降级 Dummy")
|
||
print(f"[model] 真实 joblib 不存在 ({path}), 使用 Dummy RandomForest")
|
||
return _make_dummy_model(n_features)
|
||
|
||
|
||
def _predict_block(spectra_3d: np.ndarray, model: Any) -> np.ndarray:
|
||
"""
|
||
单个 dask chunk 的推断函数 (xr.apply_ufunc 会自动调度调用)。
|
||
|
||
Parameters
|
||
----------
|
||
spectra_3d : np.ndarray
|
||
形状 (y_chunk, x_chunk, n_bands)。
|
||
此形状由 input_core_dims=[["band"]] 决定:
|
||
xarray 会把 band 维移到最后一轴, 然后按 (y, x) 的 chunk 切分调用本函数。
|
||
model : 已 fit 好的 sklearn 估计器
|
||
接受 (n_samples, n_features) 输入, 返回 (n_samples,) 预测。
|
||
|
||
Returns
|
||
-------
|
||
np.ndarray
|
||
形状 (y_chunk, x_chunk), dtype float32 的标量预测图。
|
||
"""
|
||
yc, xc, nb = spectra_3d.shape
|
||
# 2D 化: 每个像素一行光谱
|
||
flat = spectra_3d.reshape(yc * xc, nb)
|
||
# sklearn 风格的批量预测
|
||
pred = model.predict(flat)
|
||
# 还原为 2D 空间图, 强制 float32 节约一半内存
|
||
return pred.reshape(yc, xc).astype(np.float32, copy=False)
|
||
|
||
|
||
def _run_predict_sync(
|
||
model: Any,
|
||
model_id: str,
|
||
input_zarr_path: str,
|
||
output_zarr_path: str,
|
||
) -> None:
|
||
"""
|
||
同步推断主流程 (被 asyncio.to_thread 调用)。
|
||
|
||
流程:
|
||
1) xr.open_zarr 延迟打开 (dask 数组, 不一次性读入内存)
|
||
2) NaN -> 0 清洗 (model.predict 不接受 NaN)
|
||
3) xr.apply_ufunc 沿 (y, x) 逐 chunk 调 _predict_block
|
||
4) 非水域置 NaN (zarr 支持 float NaN)
|
||
5) to_zarr 触发整图计算 + 流式写出
|
||
"""
|
||
# 1. 延迟打开输入 (关键: Dask 不一次性读入内存)
|
||
ds = xr.open_zarr(input_zarr_path, chunks="auto")
|
||
if "reflectance" not in ds.data_vars:
|
||
raise KeyError(
|
||
f"输入 zarr 缺少 'reflectance' 变量; 实际: {list(ds.data_vars)}"
|
||
)
|
||
|
||
reflectance = ds["reflectance"] # dims: (y, x, band)
|
||
n_bands = reflectance.sizes["band"]
|
||
|
||
# 2. 水域掩膜 (与去耀斑算法同约定)
|
||
if "water_mask" in ds.data_vars or "water_mask" in ds.coords:
|
||
water_mask = ds["water_mask"].astype(bool)
|
||
else:
|
||
water_mask = xr.ones_like(reflectance.isel(band=0), dtype=bool)
|
||
|
||
# 3. NaN 清洗: 填充 0 (model.predict 不接受 NaN)
|
||
refl_clean = reflectance.fillna(0.0)
|
||
|
||
# 4. 核心: 用 apply_ufunc 把 model.predict 沿 (y, x) 应用
|
||
# dask="parallelized" 让每个 (y_chunk, x_chunk, band) chunk
|
||
# 独立调 _predict_block, 任意时刻内存中只有若干个 chunk。
|
||
prediction: xr.DataArray = xr.apply_ufunc(
|
||
_predict_block,
|
||
refl_clean,
|
||
kwargs={"model": model},
|
||
input_core_dims=[["band"]],
|
||
output_core_dims=[[]],
|
||
dask="parallelized",
|
||
output_dtypes=[np.float32],
|
||
dask_gufunc_kwargs={"allow_rechunk": True},
|
||
vectorize=False,
|
||
)
|
||
|
||
# 5. 非水域置 NaN (zarr 支持 float NaN, 便于后续可视化/掩膜分析)
|
||
prediction = prediction.where(water_mask, np.nan)
|
||
|
||
# 6. 包装为 Dataset 并流式写出
|
||
out = xr.Dataset(
|
||
{"prediction": prediction},
|
||
attrs={
|
||
"model_id": model_id,
|
||
"input_zarr_path": input_zarr_path,
|
||
"n_bands": n_bands,
|
||
"created_at": datetime.now().isoformat(),
|
||
},
|
||
)
|
||
# 保留 y/x 坐标
|
||
out = out.assign_coords(y=ds["y"], x=ds["x"])
|
||
|
||
# to_zarr + compute=True 触发整图 dask 图求值
|
||
# 中间会按 chunk 逐块调度到线程池, 内存峰值 ≈ 1 个 chunk 的体量
|
||
out.to_zarr(output_zarr_path, mode="w", compute=True)
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# 后台任务执行器
|
||
# ---------------------------------------------------------------------------
|
||
async def execute_train_task(
|
||
task_id: str,
|
||
model_type: str,
|
||
target: str,
|
||
train_data_path: str,
|
||
feature_start: Union[int, str],
|
||
params: Dict[str, Any],
|
||
) -> None:
|
||
"""
|
||
训练任务后台执行器（已接入真实 sklearn 训练流程）。
|
||
|
||
流程:
|
||
1) get_task 校验任务存在
|
||
2) update_task(PROCESSING)
|
||
3) 生成 model_id / model_path
|
||
4) asyncio.to_thread 派发 _run_train_sync 到默认线程池
|
||
5) 成功 -> _register_model + update_task(SUCCESS, 附 test_r2/rmse/mae)
|
||
6) 失败 -> update_task(FAILED, 附 error + traceback)
|
||
"""
|
||
record = await get_task(task_id)
|
||
if record is None:
|
||
print(f"[{task_id}] 训练任务不存在, 跳过")
|
||
return
|
||
|
||
await update_task(
|
||
task_id,
|
||
status="PROCESSING",
|
||
updated_at=datetime.now().isoformat(),
|
||
)
|
||
print(
|
||
f"[{task_id}] 开始训练 model_type={model_type} target={target} "
|
||
f"train_data_path={train_data_path} feature_start={feature_start}"
|
||
)
|
||
|
||
# model_id 用 uuid4 前 12 位 (8 位易撞, 12 位兼顾可读性)
|
||
model_id = f"model_{uuid.uuid4().hex[:12]}"
|
||
model_path = Path(f"./data/models/{model_id}.joblib")
|
||
|
||
try:
|
||
# 同步 sklearn / pandas 训练丢到默认线程池, 不阻塞 event loop
|
||
metadata = await asyncio.to_thread(
|
||
_run_train_sync,
|
||
model_type,
|
||
target,
|
||
train_data_path,
|
||
feature_start,
|
||
params,
|
||
model_path,
|
||
)
|
||
|
||
# 登记到内存注册表 (供 /predict 查 model_id)
|
||
_register_model(
|
||
{
|
||
"model_id": model_id,
|
||
"model_type": model_type,
|
||
"target": target,
|
||
"params": dict(params or {}),
|
||
"path": str(model_path),
|
||
"feature_start": feature_start,
|
||
"n_features": metadata["n_features"],
|
||
"test_r2": metadata["test_r2"],
|
||
"test_rmse": metadata["test_rmse"],
|
||
"test_mae": metadata["test_mae"],
|
||
"created_at": datetime.now().isoformat(),
|
||
"train_task_id": task_id,
|
||
}
|
||
)
|
||
|
||
# 把训练指标写回任务记录, 前端轮询时可直接看
|
||
await update_task(
|
||
task_id,
|
||
status="SUCCESS",
|
||
model_id=model_id,
|
||
model_path=str(model_path),
|
||
test_r2=metadata["test_r2"],
|
||
test_rmse=metadata["test_rmse"],
|
||
test_mae=metadata["test_mae"],
|
||
n_features=metadata["n_features"],
|
||
n_samples=metadata["n_samples"],
|
||
error=None,
|
||
traceback=None,
|
||
updated_at=datetime.now().isoformat(),
|
||
)
|
||
print(
|
||
f"[{task_id}] 训练完成 -> model_id={model_id} "
|
||
f"test_r2={metadata['test_r2']:.4f} test_rmse={metadata['test_rmse']:.4f}"
|
||
)
|
||
|
||
except Exception as exc: # noqa: BLE001
|
||
# 失败时 model_path 不一定有产物, 显式置 None 方便前端判断
|
||
await update_task(
|
||
task_id,
|
||
status="FAILED",
|
||
model_id=None,
|
||
model_path=None,
|
||
error=f"{type(exc).__name__}: {exc}",
|
||
traceback=traceback.format_exc(),
|
||
updated_at=datetime.now().isoformat(),
|
||
)
|
||
print(f"[{task_id}] 训练失败 -> {type(exc).__name__}: {exc}")
|
||
|
||
|
||
async def execute_predict_task(
|
||
task_id: str,
|
||
model_id: str,
|
||
input_zarr_path: str,
|
||
output_zarr_path: Optional[str],
|
||
) -> None:
|
||
"""
|
||
推断任务后台执行器（真实实现版）。
|
||
|
||
OOM 防护策略:
|
||
- xr.open_zarr(..., chunks="auto") 延迟打开, 整图不一次性读入内存
|
||
- xr.apply_ufunc(..., dask="parallelized") 把影像按 chunk 切分
|
||
- 每个 chunk 内部 reshape 成 2D, 调 model.predict, 再 reshape 回 2D
|
||
- 任意时刻内存峰值 ≈ 1 个 (y_chunk, x_chunk, band) chunk 的体量
|
||
- 整图完成计算后再 to_zarr(compute=True) 流式写出
|
||
"""
|
||
record = await get_task(task_id)
|
||
if record is None:
|
||
print(f"[{task_id}] 推断任务不存在, 跳过")
|
||
return
|
||
|
||
# 1. 校验 model_id 是否已注册 (避免在后台任务里报模糊错误)
|
||
model_meta = _MODEL_REGISTRY.get(model_id)
|
||
if model_meta is None:
|
||
await update_task(
|
||
task_id,
|
||
status="FAILED",
|
||
error=f"model_id 不存在: {model_id}",
|
||
updated_at=datetime.now().isoformat(),
|
||
)
|
||
print(f"[{task_id}] 推断失败 -> model_id 不存在: {model_id}")
|
||
return
|
||
|
||
# 2. 自动生成 output_zarr_path (若未提供)
|
||
if output_zarr_path is None:
|
||
stem = input_zarr_path.rstrip("/\\").split("/")[-1].split("\\")[-1]
|
||
stem = stem.replace(".zarr", "")
|
||
output_zarr_path = f"./data/{model_id}_{stem}_pred.zarr"
|
||
|
||
await update_task(
|
||
task_id,
|
||
status="PROCESSING",
|
||
updated_at=datetime.now().isoformat(),
|
||
)
|
||
print(f"[{task_id}] 开始推断 model_id={model_id} input={input_zarr_path}")
|
||
|
||
try:
|
||
# 3. 探测波段数 (用于 Dummy 模型适配)
|
||
# 这里只读 zarr 元数据 (.zarray 的 shape), 不读真实数据
|
||
ds_probe = xr.open_zarr(input_zarr_path, chunks="auto")
|
||
if "reflectance" not in ds_probe.data_vars:
|
||
raise KeyError(
|
||
f"输入 zarr 缺少 'reflectance' 变量; 实际: {list(ds_probe.data_vars)}"
|
||
)
|
||
n_bands = ds_probe["reflectance"].sizes["band"]
|
||
ds_probe.close()
|
||
|
||
# 4. 加载模型 (真实文件优先, Dummy 兜底)
|
||
model = _load_model(model_meta["path"], n_features=n_bands)
|
||
|
||
# 5. 包装同步执行, 丢到线程池, 事件循环不阻塞
|
||
await asyncio.to_thread(
|
||
_run_predict_sync,
|
||
model,
|
||
model_id,
|
||
input_zarr_path,
|
||
output_zarr_path,
|
||
)
|
||
|
||
await update_task(
|
||
task_id,
|
||
status="SUCCESS",
|
||
output_zarr_path=output_zarr_path,
|
||
model_id=model_id,
|
||
error=None,
|
||
updated_at=datetime.now().isoformat(),
|
||
)
|
||
print(f"[{task_id}] 推断完成 -> output={output_zarr_path}")
|
||
|
||
except Exception as exc: # noqa: BLE001
|
||
tb_text = traceback.format_exc()
|
||
await update_task(
|
||
task_id,
|
||
status="FAILED",
|
||
output_zarr_path=None,
|
||
error=f"{type(exc).__name__}: {exc}",
|
||
traceback=tb_text,
|
||
updated_at=datetime.now().isoformat(),
|
||
)
|
||
print(f"[{task_id}] 推断失败 -> {type(exc).__name__}: {exc}")
|
||
print(tb_text)
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# POST /api/modeling/train
|
||
# ---------------------------------------------------------------------------
|
||
@router.post("/train", response_model=TaskAcceptedResponse)
|
||
async def submit_train(
|
||
payload: TrainRequest,
|
||
background_tasks: BackgroundTasks,
|
||
) -> Dict[str, Any]:
|
||
"""提交一个模型训练任务, 立即返回 task_id。"""
|
||
|
||
task_id = str(uuid.uuid4())
|
||
await set_task(
|
||
task_id,
|
||
{
|
||
"task_id": task_id,
|
||
"kind": "train",
|
||
"model_type": payload.model_type,
|
||
"target": payload.target,
|
||
"train_data_path": payload.train_data_path,
|
||
"feature_start": payload.feature_start,
|
||
"params": payload.params,
|
||
"status": "PENDING",
|
||
"model_id": None,
|
||
"model_path": None,
|
||
"test_r2": None,
|
||
"test_rmse": None,
|
||
"test_mae": None,
|
||
"n_features": None,
|
||
"n_samples": None,
|
||
"error": None,
|
||
"traceback": None,
|
||
"created_at": datetime.now().isoformat(),
|
||
"updated_at": datetime.now().isoformat(),
|
||
},
|
||
)
|
||
background_tasks.add_task(
|
||
execute_train_task,
|
||
task_id,
|
||
payload.model_type,
|
||
payload.target,
|
||
payload.train_data_path,
|
||
payload.feature_start,
|
||
payload.params,
|
||
)
|
||
return {"task_id": task_id, "status": "PENDING", "kind": "train"}
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# GET /api/modeling/models
|
||
# ---------------------------------------------------------------------------
|
||
@router.get("/models", response_model=ModelListResponse)
|
||
async def list_trained_models() -> Dict[str, Any]:
|
||
"""
|
||
列出已训练好的模型。
|
||
|
||
未来实现: 从 ./data/models/*.joblib 扫描元信息,
|
||
当前直接从内存 _MODEL_REGISTRY 读。
|
||
"""
|
||
models = list(_MODEL_REGISTRY.values())
|
||
# 按 created_at 倒序, 最新训练的在前
|
||
models.sort(key=lambda m: m.get("created_at", ""), reverse=True)
|
||
return {"models": models, "count": len(models)}
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# POST /api/modeling/predict
|
||
# ---------------------------------------------------------------------------
|
||
@router.post("/predict", response_model=TaskAcceptedResponse)
|
||
async def submit_predict(
|
||
payload: PredictRequest,
|
||
background_tasks: BackgroundTasks,
|
||
) -> Dict[str, Any]:
|
||
"""提交一个模型推断任务, 立即返回 task_id。"""
|
||
|
||
task_id = str(uuid.uuid4())
|
||
await set_task(
|
||
task_id,
|
||
{
|
||
"task_id": task_id,
|
||
"kind": "predict",
|
||
"model_id": payload.model_id,
|
||
"input_zarr_path": payload.input_zarr_path,
|
||
"output_zarr_path": payload.output_zarr_path,
|
||
"status": "PENDING",
|
||
"error": None,
|
||
"traceback": None,
|
||
"created_at": datetime.now().isoformat(),
|
||
"updated_at": datetime.now().isoformat(),
|
||
},
|
||
)
|
||
background_tasks.add_task(
|
||
execute_predict_task,
|
||
task_id,
|
||
payload.model_id,
|
||
payload.input_zarr_path,
|
||
payload.output_zarr_path,
|
||
)
|
||
return {"task_id": task_id, "status": "PENDING", "kind": "predict"}
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# models_router — 独立于 modeling_router，路径前缀为 /models
|
||
# 最终完整路径: GET /api/models, POST /api/models/upload
|
||
# ---------------------------------------------------------------------------
|
||
models_router = APIRouter(prefix="/models", tags=["models"])
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# GET /api/models
|
||
# ---------------------------------------------------------------------------
|
||
@models_router.get("")
|
||
async def list_models() -> Dict[str, Any]:
|
||
"""
|
||
扫描 ./data/models/ 目录，返回所有 .joblib 文件名（不含后缀）。
|
||
|
||
异常处理：目录不存在时自动创建，返回空列表。
|
||
"""
|
||
models_dir = Path("./data/models")
|
||
models_dir.mkdir(parents=True, exist_ok=True)
|
||
|
||
model_names = [
|
||
p.stem for p in models_dir.iterdir() if p.suffix == ".joblib"
|
||
]
|
||
return {"models": model_names}
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# POST /api/models/upload
|
||
# ---------------------------------------------------------------------------
|
||
@models_router.post("/upload")
|
||
async def upload_model(
|
||
file: UploadFile = File(...),
|
||
) -> Dict[str, Any]:
|
||
"""
|
||
接收上传的 .joblib 模型文件，保存到 ./data/models/ 目录。
|
||
|
||
- 校验后缀必须为 .joblib
|
||
- 目录不存在时自动创建
|
||
- 返回状态和文件名（不含后缀）
|
||
"""
|
||
if not file.filename or not file.filename.lower().endswith(".joblib"):
|
||
raise HTTPException(
|
||
status_code=400,
|
||
detail="仅支持 .joblib 格式的文件",
|
||
)
|
||
|
||
models_dir = Path("./data/models")
|
||
models_dir.mkdir(parents=True, exist_ok=True)
|
||
|
||
dest_path = models_dir / file.filename
|
||
|
||
with dest_path.open("wb") as buffer:
|
||
shutil.copyfileobj(file.file, buffer)
|
||
|
||
return {
|
||
"status": "success",
|
||
"model_id": dest_path.stem,
|
||
}
|