databricks-ml-training
SkillTrain 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-trainingWhat 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 readML 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.
| Consumption | When | How |
|---|---|---|
| Batch UDF | Dashboards, 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 endpoint | Score on a user action (fraud at authorization, rec at page load) — sub-100ms | mlflow.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.
Skillssimilar to this one
All skills →dbt-agent-skills
A curated collection of Agent Skills for working with dbt, to help AI agents understand and execute dbt workflows more effectively.
640 stars
authoring-language-sdk-tasks
The language-neutral foundation for Airflow language SDKs — implement task logic in a non-Python language while the DAG stays in Python. Use when the user wants to run an Airflow task in another language (Java, Kotlin, Go, or other JVM/native languages), asks how the Python `@task.stub` pairs with n
412 stars
airflow-adapter
Airflow adapter pattern for v2/v3 API compatibility. Use when working with adapters, version detection, or adding new API methods that need to work across Airflow 2.x and 3.x.
412 stars