ClaudeSuperPower

databricks-ml-training

Skill

Train ML models on Databricks. Use for: classification/regression/deep-learning (XGBoost, scikit-learn, LightGBM, PyTorch) with Optuna, @prod/@challenger aliases, batch scoring (spark_udf for plain models, fe.score_batch for feature-store-backed), custom PyFunc, custom ResponsesAgent (LangGraph + UC

Install

git clone https://github.com/databricks/databricks-agent-skills.git ~/.claude/skills/databricks-ml-training

What is databricks-ml-training?

Train ML models on Databricks. Use for: classification/regression/deep-learning (XGBoost, scikit-learn, LightGBM, PyTorch) with Optuna, @prod/@challenger aliases, batch scoring (spark_udf for plain models, fe.score_batch for feature-store-backed), custom PyFunc, custom ResponsesAgent (LangGraph + UC Function/Vector Search); UC feature tables + FeatureLookup + point-in-time joins + Lakebase online store; declarative Feature Views (create_feature, DeltaTableSource, RollingWindow/SlidingWindow/TumblingWindow, materialize_features, streaming Kafka features). NOT for: endpoint ops (databricks-model-serving), MLflow evaluation (databricks-mlflow-evaluation).

What this can do

Capabilities declared in this component's own frontmatter — not inferred.

Inherit all session tools

Declares no tool restrictions — inherits every session tool

~165 tokens of context used while enabled, before you invoke anything

Documentation

README · ~9 min read

ML Training on Databricks

FIRST: Use the parent databricks-core skill for CLI basics, authentication, and profile selection.

Train with MLflow → register to Unity Catalog → consume the same artifact as either a batch Spark UDF over Delta or (when low-latency is required) a real-time serving endpoint.

Always train on Databricks (serverless job or notebook), never in the local Python process the agent is running in. Local training has no access to the silver tables, no MLflow tracking server, no UC registry path, and dies if the chat session drops — submit databricks jobs submit --no-wait (see "Train + deploy as a serverless job" below). Only fall back to local execution if the user explicitly asks for it.

If you need to deploy a real time model serving endpoint after the model is registered (creating endpoints, traffic config, version-swapping, querying, Foundation Model API endpoints), see databricks-model-serving.

ConsumptionWhenHow
Batch UDFDashboards, daily/hourly scores, predictions read by Genie/Dashboards or an app (often synced to a Lakebase table)mlflow.pyfunc.spark_udf(...)INSERT INTO gold_predictions. If the model was logged with fe.log_model(training_set=...), use fe.score_batch() instead — see the Feature Engineering section below.
Real-time endpointScore on a user action (fraud at authorization, rec at page load) — sub-100msmlflow.deployments.get_deploy_client() (classical) / agents.deploy() (agents). Endpoint lifecycle: see databricks-model-serving.

Default Canonical flow

Reviews

Log in to leave a review.

No reviews yet — be the first.

Explore related

Other things in this space — across every part of the ecosystem, not just skills.