databricks-ml-training

Databricks での ML トレーニング

まず最初に：CLI の基本、認証、プロファイルの選択には、親の databricks-core スキルを使用してください。

MLflow でトレーニングし、Unity Catalog に登録した後、同じアーティファクトを Delta 上のバッチ Spark UDF として、または（低レイテンシが必要な場合は）リアルタイムサービングエンドポイントとして利用します。

トレーニングは必ず Databricks 上で行ってください（サーバーレスジョブまたはノートブック）。エージェントが実行されているローカルの 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() を使用してください。以下の特徴量エンジニアリングセクションを参照してください。
リアルタイムエンドポイント	ユーザー操作時のスコアリング（承認時の不正検知、ページ読み込み時のレコメンド）— 100 ミリ秒未満	mlflow.deployments.get_deploy_client()（従来型）／agents.deploy()（エージェント）。エンドポイントのライフサイクルについては、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>) に置き換えてください。登録済みモデルの特徴量リネージに基づき、特徴量が自動的に結合されます。以下の特徴量エンジニアリングセクションを参照してください。

1 つのノートブック、1 つのアーティファクトで構成します。再実行すると再トレーニングになります。正となるデータは Gold に置きます。読み取りパスからモデルを直接呼び出すことはありません。開発中はラベルの期間ロジック（「7 日以内に障害が発生した」）をノートブック内に保持し、安定したら SDP のシルバーのマテリアライズドビューに昇格させます。

---

トレーニングと登録（大半のケース）

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 ---------- 区切り、Markdown／SQL セル用の # MAGIC %md／%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)

フレームワークの autolog：mlflow.{sklearn,xgboost,lightgbm,pytorch,tensorflow,spark}.autolog()。

ステージではなくエイリアスを使用：UC では Staging／Production が廃止されました。移動可能な @prod／@challenger を使用し、models:/{full_name}@prod で読み込みます。新しいバージョンへの昇格は、set_registered_model_alias の呼び出し 1 回で行えます。

---

利用：Delta 上でのバッチスコアリング

特徴量テーブルに基づいていないモデル向けの、低コストなデフォルトの方法です。登録済みモデルを Spark UDF として読み込み、Delta テーブルをスコアリングして、下流の利用者が読み取るゴールドテーブルに予測を書き込みます。特徴量ストアに基づくモデル（fe.log_model(training_set=...) で記録したモデル）では、spark_udf は使わず、fe.score_batch() を使ってください。詳しくは、Feature Engineering セクションを参照してください。

# 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 の 2 フィールドによる準備完了チェック、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 を参照してください。一見わかりにくい点が 2 つあります。

• 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 エージェント

手書きの ResponsesAgent（LangGraph + UC Function ツール + Vector Search による検索）については、references/genai-agents.md を参照してください。

ユーザーがカスタムの LangGraph エージェントを明示的に必要としていない限り、databricks-agent-bricks によるノーコード作成（Knowledge Assistants、Supervisor Agents）を優先してください。

---

特徴量エンジニアリング：Feature Store と Feature Views

通常の Delta テーブルではなく、どちらかの種類の特徴量エンジニアリングを使うべきなのは、同じ特徴量をトレーニング時とサービング時に同一の方法で計算する必要がある場合（トレーニングとサービングの間にずれがない）、特徴量が時間に依存し、ラベルとのポイントインタイム結合が必要な場合（未来の情報の漏洩がない）、オンラインストア経由で 10 ms 未満のレイテンシで特徴量をサービングする必要がある場合、または特徴量を複数のモデルで共有し、UC でリネージを追跡する場合です。いずれにも当てはまらない場合は、通常の UC テーブルで十分です。以下のセクションはスキップできます。

デフォルトでは、上記の理由のいずれかが明確に当てはまるか、ユーザーが明示的に求めていない限り、特徴量エンジニアリングは使わないでください。 ビルド時間と複雑さが増すためです（FeatureLookup 用に Delta テーブルの層が追加され、Feature Views 用にマテリアライゼーションパイプラインが必要になります）。判断に迷う場合は、通常の UC テーブルを使い、具体的な必要性が生じた時点で特徴量エンジニアリングを追加してください。

種類は 2 つありますが、トレーニング／スコアリングのパスは共通です（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）— 宣言した仕様（DeltaTableSource 上の create_feature、RollingWindow／SlidingWindow／TumblingWindow、Kafka の特徴量用の create_stream）に基づき、Databricks が計算、マテリアライゼーション（オフラインの Delta ＋オンラインの Lakebase）、更新を管理します。Databricks に最新の状態を保ってほしい数式としての特徴量に適しています。references/feature-views.md を参照してください。

---

注意点（時間を取られるもの）

落とし穴	対処法
モデルが UC ではなくワークスペースレジストリに登録される	ログ記録の前に mlflow.set_registry_uri("databricks-uc") を実行します
最初のクエリでエンドポイントが PERMISSION_DENIED を返す	log_model に resources=[...] を渡します（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> です

エンドポイントのライフサイクルに関する注意点（準備状態の 2 状態、バージョン切り替え、Serving UI の SP フィルター）は、databricks-model-serving に記載されています。

---

リファレンスファイル

ファイル	内容
references/custom-pyfunc.md	カスタム pyfunc の単一のエンドツーエンド例：アーティファクト、シグネチャ、code_paths、log → register → deploy → query。
references/genai-agents.md	UC Function と Vector Search ツールを使ったカスタム LangGraph ResponsesAgent。create_text_output_item の注意点と、パススルー認証用リスト resources=[...]。コード不要のエージェントには databricks-agent-bricks を推奨します。
references/feature-store.md	Unity Catalog における Feature Engineering（標準の FeatureLookup API）：FeatureEngineeringClient、timeseries_column を指定した create_table、ポイントインタイム結合を使った FeatureLookup、リネージを伴う fe.log_model、spark_udf の代わりとなる FE モデル用の fe.score_batch、および publish_table による Lakebase オンラインストア。databricks-feature-engineering>=0.16.0 が必要です。
references/feature-views.md	Feature Views（宣言型 Feature Engineering、GA）：DeltaTableSource に対する create_feature、RollingWindow／SlidingWindow／TumblingWindow 集計、materialize_features（オフライン＋オンライン）、create_stream による Kafka からのストリーミング特徴量、ポイントインタイムのトレーニングセット、および Feature Serving Endpoint。databricks-feature-engineering>=0.16.0 が必要です。

関連スキル

• databricks-model-serving — サービングエンドポイントのライフサイクル（作成、クエリ、設定更新、バージョン切り替え、AI Gateway、Foundation Model API エンドポイント）。
• databricks-agent-bricks — コード不要の Knowledge Assistants と Supervisor Agents。エージェントを手作業で構築するより、こちらを推奨します。
• databricks-mlflow-evaluation — @prod に昇格する前にモデル／エージェントの品質を評価します。
• databricks-vector-search — エージェントで検索ツールとして使用するベクトルインデックス。
• databricks-jobs — 非同期デプロイパターン（--no-wait、TASK の run_id に関する注意点）。
• databricks-unity-catalog — UC は登録済みモデルを管理します：権限、リネージ、監査。