databricks-ml-training

在 Databricks 上训练 ML

首先：使用父级 databricks-core 技能了解 CLI 基础知识、身份验证和配置文件选择。

通过 MLflow 训练 → 注册到 Unity Catalog → 将同一工件用于 Delta 上的批量 Spark UDF，或（需要低延迟时）实时服务端点。

始终在 Databricks 上训练（无服务器作业或笔记本），绝不要在 Agent 运行所在的本地 Python 进程中训练。本地训练无法访问银层表、MLflow 跟踪服务器或 UC 注册表路径，而且聊天会话中断时进程会终止——请提交 databricks jobs submit --no-wait（见下文“通过无服务器作业训练和部署”）。只有在用户明确要求时，才回退到本地执行。

如果需要在模型注册之后部署实时模型服务端点（创建端点、配置流量、切换版本、查询、Foundation Model API 端点），请参阅 databricks-model-serving。

使用方式	适用场景	方法
批量 UDF	仪表板、每日／每小时评分、供 Genie／仪表板或应用读取的预测（通常会同步到 Lakebase 表）	mlflow.pyfunc.spark_udf(...) → INSERT INTO gold_predictions。如果模型是通过 fe.log_model(training_set=...) 记录的，请改用 fe.score_batch()——见下文的 Feature Engineering 部分。
实时端点	在用户操作时进行评分（授权时进行欺诈检测、页面加载时进行推荐）——低于 100 毫秒	mlflow.deployments.get_deploy_client()（传统模型）／agents.deploy()（Agent）。端点生命周期：请参阅 databricks-model-serving。

默认标准流程

silver_<features>  +  silver_<labels>
        ▼
   notebook (as a serverless job):
   ├── train with mlflow.autolog (XGBoost / sklearn / etc.)
   ├── mlflow.register_model → UC: {catalog}.{schema}.{model}
   ├── set_registered_model_alias(name, "prod", version)
   └── spark_udf(@prod) over latest features → MERGE into gold_predictions
        ▼
gold_<entity>_predictions   ◄── dashboards, apps, Genie read this

由特征存储支持的模型在此处有所不同。 如果训练时使用了 fe.log_model(training_set=...)，请将 spark_udf(@prod) 替换为 fe.score_batch(model_uri, df=<keys_only>)——它会通过模型注册的特征血缘自动联接特征。见下文的 Feature Engineering 部分。

一个笔记本，一个工件。重新运行 = 重新训练。真实数据以 Gold 为准——读取路径绝不直接调用模型。在开发阶段，将标签窗口逻辑（failure occurred within 7 days）保留在笔记本中；稳定后，将其提升为 SDP 中的银层物化视图。

---

训练并注册（适用于 90% 的情况）

mlflow.autolog() 会为每次运行捕获参数、指标、代码和模型工件；registered_model_name=... 会自动将最佳运行注册到 UC（版本号自动递增）。使用 Optuna 包装训练，这样每次试验都是一个子运行，最佳试验将被注册。

始终使用 mlflow.set_registry_uri("databricks-uc")——否则模型会进入已弃用的工作区注册表。实验的父文件夹必须已存在——set_experiment 不会自动创建它（失败时会提示 NOT_FOUND: Parent directory does not exist）。作业运行前，使用 databricks workspace mkdirs 预先创建一次。

# Once per project — create the parent folder for the MLflow experiment.
databricks workspace mkdirs /Users/[email protected]/turbine_project

使用 Databricks 笔记本源格式（以 # Databricks notebook source 标头开头，使用 # COMMAND ---------- 分隔，并使用 # MAGIC %md／%sql 魔法命令编写 Markdown／SQL 单元格）：

# Databricks notebook source
# MAGIC %md
# MAGIC # Turbine failure prediction
# MAGIC
# MAGIC Train an XGBoost classifier on engineered turbine telemetry features.
# MAGIC ## Data exploration

# COMMAND ----------

# (basic data exploration — class balance, schema sanity, etc.)

# COMMAND ----------
# MAGIC %md
# MAGIC ## Training the model

# COMMAND ----------

import mlflow, mlflow.xgboost, optuna
from mlflow.tracking import MlflowClient
from xgboost import XGBClassifier
from sklearn.metrics import roc_auc_score

mlflow.set_registry_uri("databricks-uc")
mlflow.set_experiment("/Users/[email protected]/turbine_project/mlflow_experiment")

CATALOG, SCHEMA, NAME = "ai_demo_gen", "wind_farm", "turbine_failure"
FULL_NAME = f"{CATALOG}.{SCHEMA}.{NAME}"

# Autolog WITHOUT registered_model_name — otherwise every Optuna trial registers a new UC
# version, and a max-by-version pick lands on the last trial to finish, not the best one.
mlflow.xgboost.autolog(log_input_examples=True)

# For imbalanced labels: stratify the split, set scale_pos_weight = neg/pos.
def objective(trial):
    params = {
        "n_estimators":  trial.suggest_int("n_estimators", 100, 400),
        "max_depth":     trial.suggest_int("max_depth", 3, 10),
        "learning_rate": trial.suggest_float("learning_rate", 0.01, 0.3, log=True),
    }
    with mlflow.start_run(nested=True):
        m = XGBClassifier(**params).fit(X_train, y_train)
        return roc_auc_score(y_test, m.predict_proba(X_test)[:, 1])

with mlflow.start_run(run_name="hpo") as parent:
    study = optuna.create_study(direction="maximize")
    study.optimize(objective, n_trials=20)

# COMMAND ----------
# MAGIC %md
# MAGIC ## Retrain best params and register

# COMMAND ----------
# Retrain on the winning trial's params explicitly, then register that single model.
with mlflow.start_run(run_name="best"):
    best = XGBClassifier(**study.best_params).fit(X_train, y_train)
    mlflow.log_metric("val_auc", study.best_value)
    info = mlflow.xgboost.log_model(best, name="model", registered_model_name=FULL_NAME)

# Stages are deprecated — UC uses movable aliases. Repoint @prod at the version we just registered.
client = MlflowClient(registry_uri="databricks-uc")
client.set_registered_model_alias(FULL_NAME, "prod", info.registered_model_version)

框架自动记录：mlflow.{sklearn,xgboost,lightgbm,pytorch,tensorflow,spark}.autolog()。

使用别名，不使用阶段：UC 已移除 Staging／Production。使用可移动的 @prod／@challenger；使用 models:/{full_name}@prod 加载。提升新版本只需调用一次 set_registered_model_alias。

---

使用：在 Delta 上进行批量评分

适用于未由特征表支持的模型的低成本默认方案。将已注册模型加载为 Spark UDF，并对 Delta 表进行评分；将预测结果写入下游使用者读取的 gold 表。对于由特征存储支持的模型（使用 fe.log_model(training_set=...) 记录），请完全跳过 spark_udf，改用 fe.score_batch()——请参阅特征工程部分。

# COMMAND ----------
# MAGIC %md
# MAGIC ## Score and save to a gold predictions table

# COMMAND ----------
import mlflow
from pyspark.sql import functions as F

# env_manager rules:
#   "local"     → same runtime as training (same notebook/job). Fastest, default in dev/demo.
#   "virtualenv"→ different runtime than training; rebuilds the model's env.
#   "uv"        → same as virtualenv but faster (MLflow ≥ 2.22).
predict = mlflow.pyfunc.spark_udf(
    spark,
    model_uri=f"models:/{FULL_NAME}@prod",
    env_manager="local",
)

features = spark.table(f"{CATALOG}.{SCHEMA}.silver_turbine_features_latest")
feature_cols = [c for c in features.columns if c != "turbine_id"]   # exclude the join key
scored = features.withColumn("risk_score", predict(*[features[c] for c in feature_cols]))

# Overwrite-per-run pattern for "latest score per entity":
scored.select("turbine_id", "risk_score", F.current_timestamp().alias("scored_at")) \
    .write.mode("overwrite").saveAsTable(f"{CATALOG}.{SCHEMA}.gold_turbine_predictions")

要进行增量评分并保留历史记录，请对预测表执行 MERGE，而非覆盖。

---

实时服务（需要时）

将模型注册到 UC 后，将其部署到 Model Serving 端点。开发端的调用方式是：对于经典 ML，使用 mlflow.deployments.get_deploy_client("databricks").create_endpoint(...)；对于 ResponsesAgent，使用 agents.deploy(...)。经典 ML 的首次部署约需 5 分钟。

有关端点创建／更新／版本切换、流量配置、AI Gateway、查询、state.ready 和 state.config_update 两字段就绪检查，以及 Foundation Model API 端点的信息，请参阅 databricks-model-serving。

---

作为无服务器作业训练和部署

训练笔记本运行几分钟（Optuna + UC 注册；如果同时部署，端点预热还需 5–15 分钟）。将其提交为无服务器一次性运行，避免 CLI 阻塞。笔记本以 dbutils.notebook.exit(json.dumps({...})) 结尾，使结构化结果（model_version、val_auc、endpoint_name）能够传递到 .notebook_output.result。

# 1. Upload the training notebook
databricks workspace import /Workspace/Users/[email protected]/turbine_project/train \
  --file ./train_notebook.py --format SOURCE --language PYTHON --overwrite

# 2. Submit as serverless one-time run (returns {"run_id": N} immediately with --no-wait)
RUN_ID=$(databricks jobs submit --no-wait --json '{
  "run_name": "turbine-train-and-deploy",
  "tasks": [{
    "task_key": "train",
    "notebook_task": {"notebook_path": "/Workspace/Users/[email protected]/turbine_project/train"},
    "environment_key": "ml_env"
  }],
  "environments": [{
    "environment_key": "ml_env",
    "spec": {
      "client": "4",
      "dependencies": ["mlflow==3.1.0", "xgboost==2.1.3", "optuna==4.1.0", "scikit-learn==1.5.2"]
    }
  }]
}' | jq -r .run_id)

# 3. Poll until a terminal life_cycle_state.
for _ in $(seq 60); do
  STATE=$(databricks jobs get-run "$RUN_ID" | jq -r '.state.life_cycle_state // "UNKNOWN"')
  echo "$(date +%H:%M:%S) $STATE"
  [[ "$STATE" =~ ^(TERMINATED|SKIPPED|INTERNAL_ERROR)$ ]] && break
  sleep 30
done
[[ "$STATE" =~ ^(TERMINATED|SKIPPED|INTERNAL_ERROR)$ ]] || { databricks jobs cancel-run "$RUN_ID"; exit 1; }

# life_cycle_state TERMINATED only means "the run ended" — check result_state.
RESULT=$(databricks jobs get-run "$RUN_ID" | jq -r '.state.result_state // "UNKNOWN"')
echo "result_state=$RESULT"
[[ "$RESULT" == "SUCCESS" ]] || { echo "Run did not succeed"; exit 1; }

# 4. Pull structured output via the TASK run_id (NOT the submit run_id).
TASK_RUN_ID=$(databricks jobs get-run "$RUN_ID" | jq -r '.tasks[0].run_id')
databricks jobs get-run-output "$TASK_RUN_ID" | jq '.notebook_output.result'
# → '{"model_version":"3","val_auc":0.91,"rows_scored":124,"endpoint":"turbine-risk-endpoint"}'

请注意以下常见的 jobs submit 陷阱：无服务器笔记本任务必须设置 environments[].spec.client: "4"；get-run-output 应使用 TASK 的 run_id（tasks[0].run_id），而不是 submit 的 run_id；在无服务器一次性运行中，print() 不可靠（请使用 dbutils.notebook.exit(json.dumps(...))）；jobs submit 不接受 tags。有关更全面的 databricks-jobs 技能，请参阅 databricks-jobs。

---

自定义 pyfunc

当 sklearn/XGBoost autolog 无法满足需求时——例如需要自定义预处理、多个子模型、外部 API 调用或集成逻辑。完整示例请参阅 references/custom-pyfunc.md。有两点不太显而易见：

• python_model="path/to/file.py"（文件路径，而非类实例），并在该文件末尾调用 mlflow.models.set_model(MyModel())。这是“Models from Code”模式——文件会原样记录，不会对类进行 pickle 序列化。
• 部署前运行 mlflow.models.predict(model_uri=..., input_data=..., env_manager="uv")。这样可以在端点报错之前发现缺失的依赖。

---

自定义 GenAI Agent

手动编写 ResponsesAgent（LangGraph + UC Function 工具 + Vector Search 检索）——请参阅 references/genai-agents.md。

除非用户明确需要自定义 LangGraph Agent，否则优先通过 databricks-agent-bricks 使用无代码方式创作（Knowledge Assistants、Supervisor Agents）。

---

特征工程：Feature Store 和 Feature Views

何时应使用特征工程（任一形式），而不是普通 Delta 表：同一特征必须在训练和服务时以相同方式计算（避免训练与服务偏差）；特征依赖时间，且需要针对标签进行时间点连接（避免未来数据泄漏）；特征需要通过在线存储以低于 10 ms 的延迟提供；或者多个模型共用该特征，且需在 UC 中跟踪其血缘关系。如果这些情况都不适用，普通 UC 表就足够了——可以跳过以下部分。

默认情况下，除非上述某个理由明确适用，或用户明确提出要求，否则不要使用特征工程。 它会增加构建时间和复杂性（FeatureLookup 需要额外一层 Delta 表；Feature Views 需要物化管道）。如果不确定，请使用普通 UC 表，待出现明确需求后再添加特征工程。

两种形式，共用一套训练／评分流程（create_training_set → fe.log_model → fe.score_batch）：

• FeatureLookup API — 由你负责特征表（计算、写入、刷新）；通过 FeatureLookup 和 FeatureEngineeringClient 将它们绑定到训练集。适用于特征已存在于 Delta 表中，或你希望直接控制特征计算方式的情况。请参阅 references/feature-store.md。
• Feature Views（声明式、GA、databricks-feature-engineering>=0.16.0）— Databricks 根据你声明的规范负责计算、物化（Delta 离线 + Lakebase 在线）和刷新（通过 DeltaTableSource 上的 create_feature、RollingWindow／SlidingWindow／TumblingWindow，以及用于 Kafka 特征的 create_stream）。适用于特征是一个公式，且你希望由 Databricks 保持其最新状态的情况。请参阅 references/feature-views.md。

---

易踩的坑（那些会耗费时间的问题）

陷阱	解决方法
模型进入了工作区注册表，而不是 UC	在记录模型之前调用 mlflow.set_registry_uri("databricks-uc")
端点在首次查询时返回 PERMISSION_DENIED	将 resources=[...] 传给 log_model（涵盖 UC 函数、VS 索引、其他端点、Lakebase）——完整列表请参阅 references/genai-agents.md#resources-that-need-passthrough-auth
使用了 transition_model_version_stage	UC 中的阶段已弃用。请使用 client.set_registered_model_alias(name, "prod", version)
spark_udf 每次调用都会重建虚拟环境	在训练和评分共用同一运行时时，传入 env_manager="local"
pip_requirements 不匹配，导致端点在加载时崩溃	固定确切版本；或通过 f"mlflow=={get_distribution('mlflow').version}" 获取实时版本
agents.deploy() 生成了名称奇怪的端点	显式传入 endpoint_name=...。自动生成的名称为 agents_<catalog>-<schema>-<model>

端点生命周期相关的易踩坑点（就绪状态两态、版本切换、Serving UI 的 SP 筛选器）请参阅 databricks-model-serving。

---

参考文件

文件	内容
references/custom-pyfunc.md	单个端到端的自定义 pyfunc 示例：制品、签名、code_paths、记录日志 → 注册 → 部署 → 查询。
references/genai-agents.md	搭配 UC Function + Vector Search 工具的自定义 LangGraph ResponsesAgent。create_text_output_item 注意事项，以及 resources=[...] 直通身份验证列表。对于无代码 Agent，优先使用 databricks-agent-bricks。
references/feature-store.md	Unity Catalog 中的特征工程（标准 FeatureLookup API）：FeatureEngineeringClient、带有 timeseries_column 的 create_table、使用时间点连接的 FeatureLookup、带有血缘信息的 fe.log_model、fe.score_batch（用于 FE 模型，取代 spark_udf），以及通过 publish_table 使用 Lakebase 在线存储。需要 databricks-feature-engineering>=0.16.0。
references/feature-views.md	特征视图（声明式特征工程，GA）：对 DeltaTableSource 使用 create_feature、RollingWindow/SlidingWindow/TumblingWindow 聚合、materialize_features（离线 + 在线）、通过 create_stream 从 Kafka 获取流式特征、时间点训练集，以及特征服务端点。需要 databricks-feature-engineering>=0.16.0。

相关技能

• databricks-model-serving — 服务端点生命周期（创建、查询、update-config、版本切换、AI Gateway、Foundation Model API 端点）。
• databricks-agent-bricks — 无代码 Knowledge Assistants 和 Supervisor Agents。优先使用此项，而不是手动构建 Agent。
• databricks-mlflow-evaluation — 在将模型/Agent 晋升至 @prod 前评估其质量。
• databricks-vector-search — 用作 Agent 检索工具的向量索引。
• databricks-jobs — 异步部署模式（--no-wait、TASK run_id 陷阱）。
• databricks-unity-catalog — UC 管理已注册模型：权限、血缘信息、审计。