How can AI agents automate data science and machine learning workflows?

Summary

• AI agents can automate a large part of the ML build loop: exploring data, engineering features, training and tuning models, and running evaluations, while a data scientist frames the problem and approves what runs.
• Grounding is what makes the output usable. An agent that reads your governed tables, columns, and semantics writes code that fits your data instead of generic snippets.
• On Databricks this is Genie Code, an AI agent for data and ML work that plans a task, finds the right tables in Unity Catalog, runs code, reads the outputs to improve, and fixes its own errors.
• Work stays reproducible. Genie Code logs experiments to MLflow, so runs are tracked and comparable.
• You stay in control. The agent works within your Unity Catalog permissions and enforces existing governance while you set the direction and review its plan.

How can AI agents automate data science and machine learning workflows?

AI agents automate the build loop of data science and machine learning: the iterative work of exploring data, engineering features, training and tuning candidate models, and evaluating them. Give a capable agent a goal in plain language, such as "explore this table and build a baseline churn model," and it plans the steps, generates and runs the code, reads the results to correct itself, and keeps a reproducible record, while the data scientist sets the direction and approves what matters. The honest qualifier: these agents work best as a governed pair-programmer, not an unattended replacement, because model outcomes drive real decisions and need accuracy, reproducibility, and governance. On Databricks, this is delivered by Genie Code.

What the ML build loop involves

The build loop, often called the inner loop, is the iterative work before a model reaches production: exploratory data analysis, feature engineering, training and tuning, and evaluating candidates against your metrics. It is distinct from the operational outer loop of deployment, serving, and monitoring. Automating the inner loop means an agent can take a plain-language goal and carry out the steps end to end, generating code, running it, and iterating until the results hold up, so a practitioner works at the level of intent rather than boilerplate.

What to look for in a data science agent

Capability	Why it matters
Data grounding	The agent should read your real tables, columns, and semantics through a governed catalog, so features and code fit your data.
Multi-step planning	It should turn a goal into a plan, run code, read cell outputs, and fix its own errors instead of stopping at a first draft.
Experiment tracking	Training runs need logged parameters, metrics, and artifacts for reproducibility and comparison.
Compute awareness	The agent should select appropriate compute for a job without manual setup.
Evaluation rigor	It should run your evaluation logic on real metrics, not a fixed generic scorecard.
Permission-scoped access	It must act only within the user's existing access rights.
Human approval	You should set direction and review the plan, keeping a person accountable for what runs.

How Databricks approaches it with Genie Code

Genie Code is an AI agent built specifically for data and machine learning work. From a single prompt it plans a solution, finds relevant tables by searching Unity Catalog, runs code in a notebook, uses cell outputs to improve its results, and diagnoses and fixes its own errors. See Genie Code features and capabilities.

• Exploratory data analysis. Ask it to explore a table and surface patterns, and it builds the notebook, generates and runs the code, visualizes results, and iterates. See exploratory data analysis.
• Feature engineering grounded in your data. Genie Code reads your governed tables and semantics through Unity Catalog and reuses your team's proven patterns, so feature code stays consistent with production rather than generic.
• Model training and tuning. Ask it to train a forecasting or classification model on a specific table, and guide it toward particular model types or hyperparameter tuning. It logs experiments to MLflow and selects and configures the compute a job needs.
• Evaluation and comparison. Genie Code runs and debugs your code and compares candidate models using your own evaluation scripts and metrics.
• Governed and grounded. Integrated with Unity Catalog, Genie Code enforces existing governance and access controls, and through Lakehouse Federation it can reach data across external and on-premises systems. It can also connect to external developer tools through the open Model Context Protocol (MCP). See Introducing Genie Code and AI-assisted development.

Use cases

• From raw data to a baseline model. Ask for exploratory analysis and a first model, and the agent builds the notebook, runs the code, and iterates on the results.
• Consistent feature reuse. Generate features that follow your team's established patterns so training and production stay aligned.
• Candidate comparison. Train several variants and compare them on your evaluation metrics before you pick one to register.

FAQs

Do AI agents replace data scientists?

No. They automate repetitive building and let you work at the level of intent, but you still frame the problem, review the agent's decisions, and approve what runs and what gets registered.

What makes an agent's code trustworthy?

Grounding and governance. Genie Code reads your governed tables and semantics through Unity Catalog and acts within your permissions, so its output fits your data and respects access controls.

How is model training kept reproducible?

Genie Code logs each run to MLflow with parameters, metrics, and artifacts, so runs are tracked, comparable, and reproducible.

Which parts of the ML workflow can an agent handle?

An agent can handle the inner loop end to end: exploratory analysis, feature engineering, training and tuning, and evaluation. Deployment, serving, and monitoring belong to the operational outer loop, which is a separate stage.