ML and Generative AI in the Data Lakehouse

Building and Deploying AI Applications at Scale

With Early Release ebooks, you get books in their earliest form—the author’s raw and unedited content as they write—so you can take advantage of these technologies long before the official release of these titles.

Bennie Haelen

ML and Generative AI in the Data Lakehouse

by Bennie Haelen

Copyright © 2026 Bennie Haelen. All rights reserved.

Printed in the United States of America.

Published by O’Reilly Media, Inc., 1005 Gravenstein Highway North, Sebastopol, CA 95472.

O’Reilly books may be purchased for educational, business, or sales promotional use. Online editions are also available for most titles (http://oreilly.com). For more information, contact our corporate/institutional sales department: 800-998-9938 or corporate@oreilly.com.

Aquisitions Editor: Andy Kwan

Development Editor: Jeff Bleiel

Production Editor: Elizabeth Faerm

Copyeditor: TO COME

Proofreader: TO COME

Indexer: TO COME

Interior Designer: David Futato

Cover Designer: Karen Montgomery

Illustrator: Kate Dullea

January 2026: First Edition

Revision History for the Early Release

2024-12-11: First Release

2025-03-07: Second Release

2025-06-23: Third Release

See http://oreilly.com/catalog/errata.csp?isbn=9781098178499 for release details.

The O’Reilly logo is a registered trademark of O’Reilly Media, Inc. ML and Generative AI in the Data Lakehouse, the cover image, and related trade dress are trademarks of O’Reilly Media, Inc.

The views expressed in this work are those of the author and do not represent the publisher’s views. While the publisher and the author have used good faith efforts to ensure that the information and instructions contained in this work are accurate, the publisher and the author disclaim all responsibility for errors or omissions, including without limitation responsibility for damages resulting from the use of or reliance on this work. Use of the information and instructions contained in this work is at your own risk. If any code samples or other technology this work contains or describes is subject to open source licenses or the intellectual property rights of others, it is your responsibility to ensure that your use thereof complies with such licenses and/or rights.

978-1-098-17843-7

[LSI]

Brief Table of Contents (Not Yet Final)

Chapter 1: An Overview of Machine Learning, AI, and Generative AI (available)

Chapter 2: The Databricks Intelligence Platform (available)

Chapter 3: An Intro to ML on Databricks (available)

Chapter 4: Machine Learning Use Case with MLflow (available)

Chapter 5: AutoML (available)

Chapter 6: The Feature Store (available)

Chapter 7: Machine Learning at Scale (unavailable)

Chapter 8: Generative AI Fundamentals (unavailable)

Chapter 9: Generative AI Engineering with Databricks (unavailable)

Chapter 10: Mosaic AI on the Lakehouse (unavailable)

Chapter 11: Capstone Use Case (unavailable)

Chapter 1. An Overview of Machine Learning, AI, and Generative AI

A Note for Early Release Readers

With Early Release ebooks, you get books in their earliest form—the author’s raw and unedited content as they write—so you can take advantage of these technologies long before the official release of these titles.

This will be the 1st chapter of the final book. Please note that the GitHub repo will be made active later on.

If you have comments about how we might improve the content and/or examples in this book, or if you notice missing material within this chapter, please reach out to the editor at jbleiel@oreilly.com.

Artificial Intelligence (AI) has become deeply integrated into daily life, from virtual assistants and automated systems to sophisticated algorithms that enhance healthcare and drive financial decisions. AI has evolved from early rule-based systems to advanced machine learning (ML) and Deep Learning (DL), enabling machines to recognize patterns, learn from data, and make autonomous decisions. A new frontier, Generative AI, goes beyond analysis by creating new content, such as text, images, and music, pushing the boundaries of machine capabilities.

The Lakehouse architecture emerges as a natural fit for developing and deploying these customized solutions. Organizations are building Lakehouses that store petabytes of clean, well-organized data. While traditional BI and analytics have long derived value from this data, there’s a growing desire to leverage generative AI and large language models (LLMs) to reason over business data within the Lakehouse. This approach unlocks value at a scale that was impossible just a few years ago, enabling businesses to extract deeper insights and drive innovation through advanced AI capabilities.

This chapter will lay a foundation for the remainder of this book by providing a crisp definition of AI, ML, Deep Learning and Large Language Models.

The evolution of AI began with symbolic systems in the mid-20th century, which were limited by their inability to handle complex, real-world tasks. The 1990s and 2000s saw a shift toward data-driven approaches with the rise of ML, followed by the Deep Learning Revolution in the 2010s. This led to breakthroughs in fields like image recognition and natural language processing. By the 2020s, Generative AI emerged, enabling machines to generate creative content. Today, models like GPT-4 and OpenAI o1 use ML, deep learning, and generative AI to produce human-like language and content, driving innovations across industries.

Machine learning can be divided into sub-domains, including supervised, unsupervised, semi-supervised, and reinforcement learning. Each method serves unique purposes, from training models with labeled data to uncovering patterns in unlabeled datasets and optimizing decisions through rewards. Generative AI, which produces new content by learning patterns from data, includes models like Generative Adversarial Networks (GANs) and Variational Autoencoders (VAEs), transforming fields such as design, healthcare, and natural language processing.

While generative AI offers immense potential, it also raises ethical concerns related to data privacy, accountability, and content integrity. The reliance on personal data and the risk of creating harmful content, like deep fakes, call for strong safeguards to ensure responsible use and prevent the spread of misinformation.

In this chapter, we begin by providing clear definitions of Artificial Intelligence (AI), Machine Learning (ML), Deep Learning (DL), Large Language Models (LLMs), and Generative AI (GenAI), exploring how these concepts are interconnected. Understanding these foundational terms sets the stage for a deeper dive into the field of machine learning.

We then explore the different sub-domains of machine learning, including supervised, unsupervised, semi-supervised learning, and reinforcement learning. Each of these paradigms offers unique approaches to learning from data, and we’ll examine how they contribute to the development of intelligent systems.

Following this, we provide an overview of Generative AI, discussing the various types of generative models and highlighting typical use cases that demonstrate its transformative potential. While Generative AI holds immense promise, it also introduces significant ethical considerations. In the final section of this chapter, we delve into these concerns, addressing issues such as bias, privacy, and the implications of synthetic content generation. By acknowledging these challenges, we aim to provide a balanced perspective on the future of AI and its impact on society.

What is Artificial Intelligence?

Artificial Intelligence (AI) is the science of getting machines to accomplish tasks that match or exceed human intelligence. Intelligence is generally defined as the ability to learn, infer, reason and perform problem solving. Traditional AI systems are often built using predefined rules and logic, making them highly effective at addressing specific, well-defined problems. These systems excel in environments where the parameters are known and the tasks are predictable. However, their rigid nature presents a significant limitation: they lack the capacity to adapt or learn from new data. As a result, they struggle with the dynamic, complex tasks that require flexibility, pattern recognition, or decision-making based upon evolving inputs. This inability to grow and improve from experience constrains their effectiveness in more sophisticated applications, where adaptability is key.

Machine Learning (ML) emerged to overcome the limitations of rigid programming structures, allowing machines to learn from data patterns. Deep Learning (DL) expanded on this by enabling computers to process information in ways similar to the human brain’s neural networks. Generative AI combines principles from both ML and DL, empowering machines to create original content. Finally, Large Language Models (LLMs) are designed to process, understand, and generate human language, facilitating advanced interactions between humans and machines.

Machine Learning (ML) represents a fundamental shift in the evolution of AI, introducing the revolutionary concept of machines learning from data. Unlike traditional AI, which relies on rigid, hand-coded instructions, ML algorithms can analyze vast datasets, uncover patterns and make predictions or decisions without being explicitly programmed for every possible scenario. As a result, ML has unlocked new levels of autonomy and intelligence in AI applications.

Deep Learning, a specialized branch of Machine Learning, emulates the architecture and functioning of the human brain through the use of artificial neural networks. These networks are composed of multiple layers of interconnected nodes, or neurons, with each layer extracting increasingly complex abstract and meaningful features from the input data. This hierarchical learning process has transformed the field of AI, driving breakthrough performance in complex tasks like image and speech recognition, language translation, and autonomous driving. The explosive growth of Deep Learning has been fueled by the availability of vast datasets, often sourced from the Internet, and the GPU chips encoding the Deep Learning models. Together, these factors have propelled Deep Learning to the forefront of AI innovation.

Generative AI builds on the foundation of Deep Learning by empowering machines to go beyond analysis and prediction, allowing them to create entirely new content, whether it’s text, images, audio, music, or even video that closely mirrors human-created work. Leveraging advanced models like Generative Adversarial Networks (GANs) and Variational Autoencoders VAEs), these systems learn the underlying structure of data distributions, enabling them to produce original and lifelike outputs. This remarkable capability unlocks endless possibilities, from generating photorealistic images to creating personalized content and revolutionizing creative processes across industries.

Large Language Models (LLMs) are a powerful application of Deep Learning within the broader field of Generative AI and Machine Learning. At their core, LLMs represent a class of machine learning models designed to process and understand human language. They are trained on vast amounts of text data to identify patterns, predict the next word or phrase in a sequence, and understand complex linguistic structures. In this sense, LLMs are a specific implementation of machine learning principles, where the goal is to enable machines to learn from data without being explicitly programmed for every search.

LLMs leverage deep neural networks, composed of multiple layers of interconnected nodes (neurons). These deep learning models, particularly those built using Transformer architectures, excel at processing sequential data like language. LLMs such as GPT-4, BERT and llama have billions of parameters, allowing them to capture intricate relationships within language data. Through this deep learning framework, LLMs are able to understand context, grammar, and even subtleties of human communication. The multi-layered neural networks enable the models to generate highly sophisticated text that can be used in a variety of applications, from translation and summarization to question answering.

In the context of Generative AI, LLMs are particularly significant because they can create entirely new content based on learned patterns. This means that LLMs can do more than just analyze existing text, they can generate coherent, context-aware language outputs, such as writing essays, composing emails, or even generating code. By predicting the next word or sequence of words based on what they have learned from vast amounts of text on the Internet or book content, LLMs act as generative models, producing content that often mimics human writing.

To summarize, LLMs sit at the intersection of Machine Learning, Deep Learning and Generative AI. Machine learning provides the framework for LLMs to improve their performance through data, Deep Learning equips them with the architecture to handle complex language tasks, and generative AI showcases their ability to create new, human-like text and media. Together, these fields allow LLMs to revolutionize how machines understand and interact with language, leading to advances in a wide range of applications, from chatbots and virtual assistants to automated content generation.

The relationship between the different technologies is shown in Figure 1-1

Figure 1-1. The relationships between AI, ML, Deep Learning and Gen AI

The sub-domains of machine learning

Supervised learning involves training models with labeled datasets, where each input is paired with a known output. This method allows models to learn how to map inputs, such as images or numerical data, to the correct labels, enabling predictions on new data. Supervised learning includes classification, where outputs belong to predefined categories, and regression, which predicts continuous values. This approach is commonly used in fields like image recognition, speech translation, and sentiment analysis, relying on large labeled datasets to identify patterns in unseen data.

Unsupervised learning, on the other hand, works with unlabeled data, enabling models to detect patterns and structures without explicit guidance. Techniques like clustering and dimensionality reduction are used to reveal hidden insights. This method is essential in applications like customer segmentation, anomaly detection, and feature extraction, where labeled data is unavailable or expensive to obtain.

Semi-supervised learning combines both methods by using a small amount of labeled data alongside a larger unlabeled dataset. It helps improve accuracy in situations where obtaining labels is challenging. The model learns from labeled examples, then applies that knowledge to the unlabeled data. This technique is valuable in fields like image classification and healthcare, where labeling data is costly but vast amounts of unlabeled data exist.

Reinforcement learning involves an agent learning through interactions with its environment, receiving rewards or penalties to improve its performance over time. The agent refines its decisions by maximizing cumulative rewards, making reinforcement learning ideal for dynamic environments like robotics and autonomous driving. In models like GPT-4, reinforcement learning is used to enhance outputs, making responses more aligned with human expectations and contextually relevant.

Now let’s take a deeper look at each of these four subdomains.

Supervised Learning

In supervised learning, a model is trained on a labeled dataset where each input is paired with a corresponding label or target outcome. For example, an image classifier learns from numerous animal pictures labeled as “cat,” “dog,” etc. The goal is for the model to learn the mapping from inputs (e.g., images) to outputs (e.g., animal types) so it can accurately predict labels for new, unseen data. During training, the model adjusts its parameters to minimize the difference between its predictions and the actual labels, often using techniques like regression or classification.

A visual representation of supervised learning is shown in Figure 1-2.

Figure 1-2. Supervised learning with labeled dataset

Classification and regression are two fundamental types of supervised learning in machine learning, where the model learns from labeled data to make predictions.

The goal of classification is to predict the category or class of an input based on its features. In classification, the output is discrete, meaning it belongs to one of a predefined set of categories. An binary example of classification would be a when there are only two possible outcomes, for example an email is either spam or it is not, while a multiclass example would be a case where there are more than two possible classes, and is the case in our animal classification above, where the class is either “cat”, “dog” or “bear”.

Regression is used to predict a continuous numerical value based on the input features. Instead of discrete categories, the predicted output in regression is a real number. Examples include predicting home prices based on location, number of bedrooms and bathrooms etc. or time series forecasting where we predict future stock prices over time.

Supervised learning is often used for use cases where large amounts of historical data, together with their outcome is available. Common examples include:

Image recognition

Identifying objects, like the animal example above, recognizing faces or scenes in images. Large amounts of images with their captions are available in the public domain, including social media sites.

Speech recognition and language translation

The ability to translate spoken language into text, and to translate one language to another. Again, a vast volume of audio and text media is available to train these models.

Loan approval

Loan application approvals, banks have large amounts of past loan applications, together with their final outcome (“approved” or “denied”).

Sentiment analysis

Determining the sentiment (positive, negative or neutral) of text data or content. Large amounts of social media data is available to train these models.

Unsupervised Learning

In unsupervised learning, a model is trained on unlabeled data without explicit instructions. The model examines the underlying structure to detect patterns, relationships, and groupings, uncovering hidden insights that might not be immediately apparent. This approach is particularly useful when labeled data is unavailable or when exploring the dataset’s inherent structure.

Key techniques include:

Clustering

Grouping similar data points based on shared characteristics to reveal natural groupings.

Association

Identifying relationships or dependencies between data items to discover how certain attributes influence each other.

Dimensionality Reduction

Simplifying the dataset by reducing the number of variables or features while preserving core information, making it easier to visualize complex data and enhancing model performance by removing irrelevant or redundant data.

An overview of unsupervised learning is provided in Figure 1-3. Unsupervised learning is typically performed in three steps:

1. Data Exploration

This step involves analyzing and understanding the raw data, identifying the features, distributions, and patterns that can guide the unsupervised learning process. It sets the foundation by revealing potential insights without the guidance of labeled data.

2. Model Selection

Here, the focus is on choosing the most suitable unsupervised learning model or technique for the task at hand. This includes selecting from methods like clustering (e.g., K-Means), association (e.g., Apriori), or dimensionality reduction (e.g., PCA).

3. Pattern Discovery

This step involves applying the chosen algorithm to extract meaningful patterns, relationships, or clusters from the data. It’s the core of unsupervised learning, where insights are gained and actionable groupings or associations are identified.

Figure 1-3. Unsupervised learning stages

Unsupervised learning techniques are applied in a wide range of use cases. For example, in customer segmentation, clustering algorithms are used to group customers based on behaviors or demographic information, allowing businesses to tailor marketing strategies. Anomaly detection is another critical application, where the model identifies unusual patterns in data, such as fraudulent transactions in credit card usage or abnormal network activity that may indicate a security breach. Dimensionality reduction, on the other hand, is often used in tasks like feature extraction, where new, more relevant features are generated from the original data, and noise reduction, which helps to eliminate irrelevant signals and background noise from datasets, making the core information clearer and more actionable.

These examples highlight how clustering, association, and dimensionality reduction allow machine learning models to discover hidden insights in data, making unsupervised learning a powerful tool in scenarios where labeled data is unavailable or impractical to generate.

Semi-Supervised Learning

Semi-supervised learning is a hybrid machine learning approach where a small amount of labeled data is combined with a large amount of unlabeled data during model training. This allows the model to learn from vast unlabeled datasets while being guided by labeled examples, making it useful when labeling data is difficult or expensive. Initially, the model learns from the labeled data to gain a basic understanding of the problem domain. It then incorporates the unlabeled data, extending its knowledge by discovering patterns and structures within it.

For example, consider images of cats, dogs, and bears where only the cat and dog images are labeled. The model first learns to identify cats and dogs from the labeled data. It then recognizes that some images don’t fit these categories (the bears) and, upon minimal additional labeling (identifying these as bears), the model improves its ability to classify all three animals accurately, as shown in Figure 1-4.

Figure 1-4. Unsupervised Models with partially labeled data

Semi-supervised learning is particularly advantageous in scenarios where labeling large volumes of data is costly, time-consuming, or otherwise impractical. As mentioned earlier, this learning method strikes a balance between fully supervised and unsupervised approaches by utilizing a small amount of labeled data in conjunction with a much larger set of unlabeled data. Several real-world applications illustrate the power of semi-supervised learning across diverse domains.

In image and video analysis, where manually labeling media files is labor-intensive, semi-supervised models can classify images or video sequences using limited labeled samples. These models leverage the structure and patterns in the vast pool of unlabeled media to improve accuracy and reduce the need for extensive labeling.

Healthcare is another critical domain for semi-supervised learning, especially in areas such as disease prediction and patient outcome analysis. Medical datasets often contain a small proportion of labeled patient records, such as those detailing the onset of a disease, while the majority remain unlabeled. Semi-supervised models can utilize the smaller labeled set to predict outcomes more accurately across the larger dataset, improving medical diagnostics and patient care.

In e-commerce, semi-supervised learning is applied to product categorization. Retailers often have labeled data for specific product categories, but the majority of products may not be fully classified. By combining labeled and unlabeled product data, semi-supervised models enhance the ability to categorize products more efficiently, improving search, recommendation systems, and inventory management.

Reinforcement learning

To introduce reinforcement learning, consider training a dog: you instruct the dog to perform tasks and reward it with treats when it succeeds. The dog learns through positive feedback to repeat desired behaviors. Similarly, negative feedback, like scolding, discourages unwanted actions.

In reinforcement learning, an agent learns by interacting with an environment, receiving rewards or penalties based on its actions. The goal is to develop a policy that maximizes cumulative rewards over time through trial and error. The agent starts in a state, takes an action, and the environment provides a new state along with a reward or penalty. Over time, the agent associates actions with outcomes and improves its decision-making process. Figure 1-5 illustrates this feedback loop between the agent and the environment, where optimal actions are reinforced.

Figure 1-5. The reinforcement learning feedback loop

The value of Reinforcement Learning (RL) lies in its ability to solve highly complex problems by mimicking human learning processes. The model undergoes continuous learning and improvement through interaction with its environment, making it highly accurate over time. Reinforcement learning is especially valuable when applied to tasks like robotics, autonomous driving, and simulators, where it constantly learns from new experiences. It also integrates well with neural networks, creating the foundation for deep reinforcement learning, which allows the model to handle even more intricate tasks. RL’s adaptability and experience-based learning make it a powerful tool for solving complex problems where traditional models may struggle.

However, Reinforcement Learning also comes with significant challenges. It requires vast amounts of training data and computational resources, which makes it inefficient for simpler tasks where less complex models would suffice. The training process can be time-consuming and expensive, both in terms of computation power and maintenance costs, especially for real-world applications like autonomous vehicles and robotics. Additionally, excessive training may lead to issues such as overfitting, where the model becomes too focused on its past experiences, potentially affecting its ability to generalize to new situations. These drawbacks limit RL’s practical use in simpler or resource-constrained scenarios.

Reinforcement learning is used with models like GPT-4, particularly through a technique called Reinforcement Learning from Human Feedback ¹ (RLHF). This approach involves training the model to generate responses that align more closely with human preferences, improving the quality and coherence of its outputs. Like earlier models, GPT-4 starts with unsupervised learning, where it is trained on vast amounts of text data to predict the next word in a sequence. This stage does not involve reinforcement learning.

In the next stage RLHF is applied as follows:

Human annotators rank the model’s responses based on quality and relevance.

These rankings are used as a reward signal in a reinforcement learning loop.

The model is then fine-tuned to generate responses that maximize this reward, aligning more closely with human expectations.

This reinforcement learning step helps improve GPT-4’s ability to generate responses that are not only accurate but also contextually appropriate, polite, and engaging.

By combining unsupervised learning with reinforcement learning, models like GPT-4 are able to generate more sophisticated and human-like interactions, making them more useful across a variety of applications, from chatbots to creative writing tools.

An overview of Generative AI

Generative AI is a type of artificial intelligence that can generate new content, such as text, images, music, or video. It is a class of artificial intelligence models that are designed to generate new data or content that resembles existing data. Traditional AI models tend to be primarily focused on classification, regression or decision-making, but Generative AI models learn the underlying patterns and structure of the input data and use this knowledge to create new, synthetic instances that are similar to the original data.

The result could include generating realistic images, writing coherent text, composing music, designing products, or even simulating entire environments. This is done by training the models on a dataset, where they learn the underlying patterns, distributions and relationships within the data. After training, these models can produce new data that looks similar to the training data but is not an exact copy. For example, a generative model trained on a dataset of human faces can generate new, unique faces that don’t belong to any real person.

Types of Generative Models

Generative AI models include Generative Adversarial Networks (GANs), which pit two neural networks against each other to produce realistic data; Variational Autoencoders (VAEs), which encode input data into a latent space and decode it to generate new samples; and transformer-based models, like GPT-4, which leverage attention mechanisms to generate coherent and contextually relevant sequences.

Generative Adversarial Networks

Generative Adversarial Networks (GANs) have been utilized effectively for a number of applications, such as text generation, music composition, and picture synthesis.

GANs consist of two neural networks: a generator and a discriminator which are competing against each other. The generator tries to create realistic (but fake) data, while the discriminator tries to distinguish between real and generated data. This competition leads to the generator producing increasingly realistic outputs.GANs are commonly used in image and video generation tasks, where they have shown impressive results in generating realistic images, creating animations, and even generating synthetic human faces. An example of a GAN is shown in Figure 1-6

Figure 1-6. Generative Adversarial Network Architecture

Variational Autoencoders

Variational autoencoders, or VAEs, are a class of generative models known for their ability to uncover hidden representations within data. These models consist of two neural networks, commonly referred to as the encoder and decoder. The role of the encoder is to take a sample of the input data and transform it into a latent representation which is a compact numerical vector that encapsulates the key characteristics of the data. Thus, if the model generates faces, the latent space will contain features, such as gender, facial expression, age, and more.

Note

A latent representation refers to a compressed or abstract version of the input data that captures its most essential features in a lower-dimensional space. It is used in many machine learning models, particularly in autoencoders, variational autoencoders (VAEs), and deep learning models for tasks like dimensionality reduction, generative modeling, and feature learning.

The decoder, in turn, utilizes this latent representation to generate a reconstruction of the original input. This process allows VAEs to effectively learn meaningful patterns within the data, facilitating tasks such as generating new samples or compressing information into a lower-dimensional space. VAEs are powerful tools for generating realistic human faces and images, text generation, music composition, treatment findings, and drug discovery. A model of a VAE is shown in Figure 1-7

Figure 1-7. The Variable Autoencoder Model

Transformer-based Models

Transformer-based models have demonstrated remarkable capabilities in handling natural language processing (NLP) tasks, such as language translation and question answering. At their core, these models are built on a unique architecture known as the transformer, which consists of two key components: an encoder and a decoder. The encoder is responsible for analyzing the input data, while the decoder uses the encoded information to generate a corresponding output. What makes transformers particularly powerful is their ability to learn from large amounts of unlabeled data, enabling them to capture complex contextual relationships within sequential data. A high-level representation of a transformer model (in this case, one which is intended to translate from one language to another) is shown in Figure 1-8.

Figure 1-8. The Transformer architecture

A key advantage of transformer models lies in their capability to handle multiple sequences simultaneously, as opposed to processing them sequentially. This parallelism allows transformers to efficiently manage large datasets, making them particularly suited for complex tasks. The model’s ability to understand language relies on two important techniques: self-attention and positional encoding.

Self-attention² plays a pivotal role by enabling the model to discern the context of each word in a sequence, which is essential for understanding the full meaning of a sentence. Positional encoding³, on the other hand, addresses the model’s inherent lack of awareness regarding word order, providing crucial information about the relative positions of words within a sequence. Together, these mechanisms enable transformers to capture the nuances of language.

These architectural innovations have been instrumental in the rise of models like GPT (Generative Pre-trained Transformers), which are now widely applied in text generation. By harnessing the strengths of transformer-based systems, GPT models have expanded the potential for automating content creation and producing text that closely resembles human communication.

Generative AI Applications and Use Cases

Generative AI is having a profound impact across many industries and professions, from creative fields like arts and entertainment to sectors like healthcare and science. These models can handle repetitive tasks, offer personalized experiences, and tackle complex challenges in ways that streamline workflows and improve outcomes.

One area where generative AI is making a mark is art and design.

In medicine and healthcare, generative AI is transforming how professionals diagnose illnesses, plan treatments, and develop new medications. The models are also proving useful in analyzing medical images, contributing to more precise and personalized treatments. By automating routine processes, healthcare providers can reduce costs and improve the overall quality of care. These AI systems are helping to reshape the industry, offering more targeted solutions that improve patient outcomes and streamline operations.

When it comes to Natural Language Processing (NLP), generative AI models are enabling computers to produce human-like language, which is key for chatbots, virtual assistants, and content creation tools. They are also used for tasks like language modeling, sentiment analysis, and text summarization. These capabilities help organizations manage large amounts of textual data, automate customer service, and improve the efficiency of content production. By enhancing the way machines understand and respond to human language, generative AI is driving advancements in NLP that could change how we interact with technology on a daily basis.

Ethical Considerations

As generative AI continues to evolve and become more integrated into various industries, it brings with it a range of powerful capabilities, from content creation to automation. However, these advancements also introduce important ethical considerations that must be carefully addressed. The potential for misuse, unintended consequences, and the need for responsible handling of AI-generated content highlights the importance of developing and deploying these technologies with strong ethical guidelines.

If ethical lapses or AI failures occur, the impact on a business can be significant. Misinformation, biases, or harmful content generated by AI can lead to reputational damage, customer distrust, and potential regulatory scrutiny. The public relations fallout from an AI-driven error or ethical misstep can erode consumer confidence, resulting in lost revenue and lasting harm to brand image. Businesses, therefore, need to proactively address ethical considerations in AI implementation, not only to ensure compliance but also to protect and strengthen their reputation in a highly competitive, and increasingly transparent, marketplace.

We will now explore several key areas where ethical considerations are crucial to ensuring that generative AI is used responsibly and safely. Each of these areas is vital to ensuring that AI technologies are used in a way that benefits society while minimizing risks and potential harm.

Data Privacy

Generative AI models typically require vast datasets (which often include sensitive or personal information such as medical records, social media posts, or financial data) in order to function effectively. This creates significant privacy concerns, particularly if the data is collected or used without informed consent. Organizations must implement strong data governance practices to ensure that all personal data is anonymized, securely stored, and handled according to legal regulations, such as GDPR or HIPAA. Ethical AI systems should prioritize user privacy and seek to minimize the collection of sensitive data wherever possible.

Accountability

As AI systems become more autonomous and integrated into decision-making processes, the question of accountability grows in importance. Businesses and developers must ensure that AI systems operate within predefined ethical boundaries and that there are mechanisms to identify and correct issues when things go wrong. Since AI models can make decisions that directly impact people’s lives—such as credit approvals, hiring decisions, or healthcare recommendations—organizations must maintain accountability for these outcomes. Clear guidelines should be in place to determine who is responsible when an AI system makes a mistake or causes harm.

Transparency

Transparency is crucial for maintaining trust in AI systems. When generative AI models are used to produce content or make decisions, it is important to communicate clearly how these models contribute to the process. This means explaining what the AI is doing, how it’s being used, and the potential biases or limitations inherent in its design. For instance, if generative AI is used in medical diagnosis, patients should be informed about its role in their care, including the strengths and weaknesses of the AI’s decision-making abilities. Transparency also involves disclosing how the data used to train the AI was sourced, allowing users to understand and assess its reliability.

Explainability

Explainability refers to the ability to understand and articulate how AI models make decisions or generate outputs. Many AI models, particularly deep learning models, operate as “black boxes,” making it difficult to interpret their inner workings. This lack of clarity can raise ethical concerns, especially when AI decisions impact critical areas such as healthcare, finance, or law enforcement. It is essential to develop models that can be explained in understandable terms, ensuring that users and stakeholders can trust that the AI is making decisions based on clear, logical processes. Explainability is key to ensuring fairness, safety, and ethical use of AI.

Privacy Risks

Generative AI has the ability to create highly realistic images, videos, and text, which can be easily manipulated or exploited for unethical purposes. For example, AI-generated content could be used for identity theft, deepfakes, or fraudulent activities, where realistic but false content is used to deceive others. Such risks emphasize the need for strong safeguards, including the development of detection tools to identify AI-generated forgeries. Organizations must be mindful of these risks and implement policies that prevent the use of generative AI in ways that could harm individuals or society at large.

Societal Impact

The widespread dissemination of AI-generated content can have far-reaching effects on society, some of which may be unintended. The ability to generate realistic content at scale could contribute to misinformation, influence public opinion, or lead to reputational damage for individuals or organizations. If not managed responsibly, generative AI could exacerbate issues related to trust in media, politics, and institutions. Businesses using these technologies must consider the broader societal implications and work proactively to mitigate harmful consequences, ensuring that AI-generated content does not cause harm to the public or undermine societal values.

Content Integrity

Generative AI can produce virtually anything based on input prompts, which raises concerns about the integrity of the content it generates. Ensuring content integrity requires developing ethical guidelines for the creation and use of AI-generated content. For example, AI should not be used to generate malicious deep fakes, discriminatory content, or misinformation. Additionally, content moderation and filtering mechanisms should be in place to prevent the spread of harmful AI-generated materials, ensuring that generative AI is used responsibly and ethically.

Summary

In this chapter, you gained a comprehensive overview of Artificial Intelligence (AI), Machine Learning (ML), Deep Learning (DL), Large Language Models (LLMs), and Generative AI, establishing a foundation for understanding these interconnected concepts. You explored the evolution of AI from early rule-based systems to modern ML and DL approaches that enable machines to learn from data and mimic human cognitive functions. The chapter delved into the sub-domains of ML, including supervised, unsupervised, semi-supervised, and reinforcement learning, explaining how each method contributes to building intelligent systems. You also examined Generative AI, highlighting types of generative models such as Generative Adversarial Networks (GANs), Variational Autoencoders (VAEs), and transformer-based models like GPT-4, and investigated their applications across various industries.

Furthermore, you considered the ethical considerations surrounding Generative AI, emphasizing the importance of data privacy, accountability, transparency, explainability, and content integrity. You acknowledged the potential risks associated with AI-generated content, such as misinformation and privacy violations, and understood the need for responsible development and deployment of AI technologies to minimize harm and maximize societal benefits.

In the next chapter, you will be introduced to the Databricks Intelligence Platform, a unified data analytics and AI solution that empowers organizations to build and deploy machine learning models efficiently. Building on the concepts discussed here, you’ll explore how Databricks facilitates the end-to-end machine learning lifecycle within a lakehouse environment, enabling you to harness the full potential of AI and Generative AI in practical applications.

¹ What is reinforcement learning from human feedback (RLHF) by Dave Bergmann. IBM, November 2023

² Attention is all you need, by Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N. Gomez, Lukasz Kaiser, Illia Polosukhin, June 2017, Computer Science Journal

³ What is positional encoding in the transformer model? Data Science

Chapter 2. The Databricks Intelligence Platform

A Note for Early Release Readers

With Early Release ebooks, you get books in their earliest form—the author’s raw and unedited content as they write—so you can take advantage of these technologies long before the official release of these titles.

This will be the 2nd chapter of the final book. Please note that the GitHub repo will be made active later on.

If you have comments about how we might improve the content and/or examples in this book, or if you notice missing material within this chapter, please reach out to the editor at jbleiel@oreilly.com.

This chapter explores the reasons for adopting the Databricks Platform as the platform of choice for developing advanced AI, machine learning (ML), and generative AI solutions.

This chapter provides an extensive overview of the Databricks Platform, highlighting its key components such as the optimized runtime environment, versatile cluster configurations, collaborative workspaces, and interactive notebooks. The Databricks Runtime enhances performance through proprietary optimizations like Photon Acceleration, ensuring efficient data processing and rapid model training. Clusters are categorized into All-Purpose and Job Clusters, each tailored to support interactive development and automated workflows respectively. Workspaces facilitate organized collaboration by allowing teams to manage notebooks, libraries, and resources within a secure and scalable environment.

In this chapter, we’ll also walk you through all the steps needed to set up a Databricks Machine Learning environment. You’ll learn how to create a Data account and workspace, how to configure that workspace, how to set up clusters, and how to select a runtime.

Why the Databricks Platform?

The Databricks platform is well-suited for AI, machine learning (ML), and generative AI development due to its seamless integration of data engineering, data science, and ML workflows within a unified Lakehouse architecture eliminating data silos and facilitating collaboration across teams. By combining the scalability and flexibility of data lakes with the performance and reliability of data warehouses, the Databricks Lakehouse provides a single environment where vast amounts of structured and unstructured data can be efficiently stored, processed, and analyzed. Its scalable, high-performance processing engine powered by Apache Spark, and its comprehensive security features make it an ideal choice for organizations aiming to leverage large-scale data for innovative AI-driven initiatives.

The Databricks platform has a comprehensive suite of collaborative tools and managed infrastructure, which empower teams to innovate without the overhead of managing complex environments. Interactive notebooks, version control integrations, and real-time collaboration features facilitate a productive workflow for data scientists and ML engineers, allowing them to prototype, iterate, and refine models seamlessly. The platform’s support for a wide range of ML frameworks and libraries, including TensorFlow, PyTorch, and Scikit-learn, ensures that developers have the flexibility to leverage the best tools for their specific needs. Databricks’ managed clusters, featuring auto-scaling and optimized runtimes, provide the computational power necessary for training sophisticated generative AI models, while also ensuring cost-efficiency and operational reliability. Additionally, Databricks offers a suite of features specifically designed for machine learning, including integrated MLflow for model lifecycle management, support for popular ML frameworks, and advanced tools for distributed training and hyperparameter tuning. These features collectively create a powerful ecosystem that supports the end-to-end development and deployment of sophisticated AI and ML models.

The Databricks platform offers robust security features and comprehensive governance controls. The platform incorporates enterprise-grade security measures, including data encryption, role-based access control (RBAC), and compliance with industry standards such as GDPR and HIPAA. These features ensure that sensitive data remains protected throughout its lifecycle, from ingestion to model deployment. Databricks’ integrated monitoring and logging tools provide deep visibility into system performance and user activities, enabling proactive management and quick resolution of issues. This strong emphasis on security and governance not only safeguards valuable data assets but also builds trust among stakeholders.

Databricks Platform Overview

The Databricks Data Intelligence Platform is an integrated solution designed to unify data engineering, data science, machine learning, and business analytics within a single ecosystem.

An overview of the Databricks Intelligence Platform is shown in Figure 2-1.

Figure 2-1. The Databricks Data Intelligence Platform

The Open Data Lake at the foundation of the platform serves as a centralized repository where all raw data is stored and managed. It is designed to handle vast amounts of diverse data types, including data, logs, text, audio and images and other structured or unstructured data. The Open Data Lake plays a pivotal role in enabling data engineering. It offers a flexible, scalable and highly accessible storage layer for all data, irrespective of its format or origin.

Delta Lake is an open-source storage layer that brings reliability, performance and scalability to data lakes. Delta Lake enhances traditional data lakes by adding key features such as ACID transactions, versioning and schema enforcement, all of which are critical for managing large-scale operations and supporting advanced analytics and machine learning workflows.

Unity Catalog provides comprehensive data governance, security and accessibility across the entire data intelligence platform. By centralizing the management of access controls, lineage, data discovery and compliance, Unity Catalog ensures that organizations can safely and efficiently manage their data assets at scale. It plays a foundational role in enabling secure, governed data collaboration both within the organization and with external stakeholders supporting the growing complexity of modern data ecosystems.

The Workflows component plays a crucial role in automating and orchestrating complex data and ML pipelines. Workflows manage the end-to-end execution of tasks, such as data transformations, model training, and reporting, in a structured and efficient manner. By using job cost optimization based on past runs, it ensures that workflows are executed efficiently in terms of both performance and resource usage, thereby reducing operational costs.

Delta Live Tables are designed to simplify and automate the process of building and managing reliable, high-quality data pipelines. It enhances data engineering workflows by automating tasks such as data validation, transformation, and quality checks, ensuring that data remains consistent and up to date as it moves through the pipeline. Delta Live Tables integrates tightly with Delta Lake, providing real-time data processing capabilities and ensuring that all data pipelines are optimized and reliable.

Databricks SQL enables users to query and analyze data using SQL, while also providing advanced features such as text-to-SQL capabilities. It is designed to make data accessible and actionable for a wide range of users, including data engineers, analysts, and non-technical stakeholders. Databricks SQL simplifies the querying process, allowing users to interact with large datasets in the Open Data Lake and Delta Lake without needing to understand the underlying infrastructure.

Mosaic AI provides tools for building, fine-tuning, and deploying custom ML models. It plays a crucial role in empowering data scientists and ML engineers to develop sophisticated AI models that are highly customized for specific use cases. Mosaic AI is designed to facilitate the entire machine learning lifecycle, from model development to real-time deployment and monitoring.

The Databricks Runtime

This section provides an in-depth overview of the Databricks Runtime environment, exploring its architecture, features, and how it facilitates advanced analytics and AI workloads in a unified data platform.

The Databricks runtime is a data processing engine built on a highly optimized version of Apache Spark, tailored to accelerate analytics, data science, and machine learning workloads. It serves as the execution environment in Databricks clusters, providing a robust foundation for running large-scale data processing tasks with improved performance and reliability.

At its core, Databricks Runtime integrates a multitude of optimizations and proprietary enhancements that extend beyond standard Spark capabilities. These enhancements include performance optimizations, security features, and compatibility with a wide array of data sources and libraries essential for ML and AI applications.

Key Components of the Databricks Runtime

The Databricks Runtime consists of several key components that together provide a powerful environment for data processing, analytics, and machine learning.

Optimized Spark Engine

At the core of the Databricks Runtime is its optimized Apache Spark engine, which includes numerous enhancements to improve performance and efficiency. One of the significant advancements is the Photon Engine, a vectorized query engine designed to significantly accelerate SQL and Spark workloads by leveraging modern CPU architectures. This engine enables faster data processing and query execution, reducing latency in analytics tasks.

Another vital component is the Delta Engine, which optimizes data lake storage through Delta Lake. It introduces ACID transactions, allowing for reliable and consistent data operations, and provides scalable metadata handling and data versioning. These features ensure data integrity and simplify data data management in large-scale environments. Additionally, the runtime incorporates improved Delta Caching which reduces I/O operations, thereby accelerating data access and processing.

Integrated Machine Learning Libraries

The Databricks Runtime comes pre-installed with a rich set of machine learning libraries and frameworks, facilitating the development of advanced analytics and AI solutions. It includes MLLib, Spark’s scalable machine learning library that supports various algorithms for classification, regression, clustering and collaborative filtering. This library enables data scientists to build models that can handle large datasets efficiently.

For deep learning applications, the runtime provides support for popular frameworks such as TensorFlow, PyTorch, and Keras, allowing users to build and train complex neural networks. It also includes XGBoost and LightGBM, which are powerful gradient boosting algorithms known for their performance in predictive modeling tasks. To manage the machine learning lifecycle effectively, MLFlow is seamlessly integrated into the environment, offering tools for experiment tracking, model management, and deployment.

Language Support

The Databricks Runtime offers support for multiple programming languages. Python is fully supported, including PySpark for distributed processing and access to popular data science libraries, like pandas, NumPy, and SciPy. This makes it accessible for data scientists who are familiar with Python’s rich ecosystem.

For those who prefer functional programming paradigms, the runtime environment also supports Scala, leveraging Spark’s native language for high-performance data processing tasks. The environment also supports SQL, enabling traditional database querying capabilities for data analysis and manipulation. Furthermore, R is supported for statistical computing and graphics, integrating with SparkR and sparklyr packages, which allows statisticians and data analysts to utilize their existing skill sets within the Spark framework.

Advanced Analytics and AI Capabilities

The Databricks Runtime extends beyond basic data processing by providing advanced analytics and AI capabilities. It supports graph processing through GraphFrames, enabling users to perform graph analytics essential for applications like social network analysis and fraud detection. This allows for the exploration of relationships and patterns within connected data.

The runtime also offers robust support for streaming analytics via Spark Structured Streaming, facilitating real-time data processing. This is crucial for applications that require immediate insights from continuously flowing data, such as monitoring systems and live dashboards. For natural language processing (NLP) tasks, the environment includes libraries and tools for text analysis, sentiment analysis, and other NLP functionalities, empowering users to derive insights from unstructured textual data.

Importantly, the Databricks Runtime is equipped to handle the development and deployment of Generative AI models. Leveraging large-scale data processing capabilities, users can train and implement generative models that are foundational for applications like text generation, summarization, and conversational AI.

Runtime Summary

By integrating an optimized Spark Engine, comprehensive machine learning libraries, multi-language support, and advanced analytics capabilities, the Databricks Runtime environment empowers organizations to execute complex data workloads efficiently. This holistic approach facilitates the development of ML, AI, and Generative AI solutions within the Lakehouse architecture, driving innovation and enabling data-driven decision making across enterprises.

Clusters, Workspaces, and Notebooks

In our exploration of AI, ML and Generative AI within a Lakehouse environment, understanding the foundational tools and components is essential. This section delves into the core elements of Databricks, which are Clusters, Workspaces and Notebooks, and how these components work together to facilitate advanced analytics and AI development in a Lakehouse architecture.

Understanding Databricks Clusters

In a Lakehouse environment, clusters play a pivotal role by providing the computational resources necessary to process data stored in the data lake and perform complex analytics tasks. They enable data ingestion and transformation, allowing organizations to process raw data into structured formats suitable for analysis and machine learning. Clusters facilitate machine learning model training by offering scalable compute power to handle large datasets and complex algorithms. Additionally, they support real-time analytics through streaming data processing capabilities, providing up-to-date insights for decision-making.

Understanding how to effectively utilize clusters requires a grasp of the typical development workflows within Databricks. These workflows are designed to transition efficiently from data ingestion to model deployment and can be broadly categorized into two distinct types: Interactive Development and Automated Execution. Mastery of these workflows is essential for building robust AI, machine learning (ML), and generative AI solutions.

Interactive Development Workflow

The interactive development workflow is a hands-on, iterative process aimed at achieving specific objectives, such as ingesting data from various sources and transforming it into a feature set suitable for training machine learning models. This workflow typically involves the following steps:

1. Defining Objectives

Initiate the development cycle by outlining clear goals, such as data ingestion, preprocessing, feature engineering, and model training.

2. Notebook-Based Development

Launch a new Databricks notebook, an interactive environment for writing and executing code. Notebooks support multiple languages, including Python, Scala, R, and SQL, enabling flexible and comprehensive data manipulation and analysis.

3. Interactive Coding and Execution

Enter and execute commands within individual cells of the notebook. This cell-based execution facilitates an iterative approach, allowing developers to test hypotheses, visualize data, and refine algorithms in real time.

4. Continuous Refinement

Iterate through the development process, making adjustments based on intermediate results and insights until the desired outcomes are achieved.

This interactive workflow is supported by All-Purpose Clusters in Databricks. These clusters are optimized for interactive use, allowing developers to execute notebook commands efficiently during the development phase. All-Purpose Clusters provide the computational resources necessary to handle complex data transformations and model training tasks in an exploratory manner. Additionally, multiple users can share All-Purpose Clusters to perform collaborative interactive analysis, with users possessing the appropriate permissions able to terminate, restart, and detach clusters from multiple notebooks.

Automated Execution Workflow

Upon completing the development and testing phases, the next step involves transitioning the code from an interactive environment to an automated execution workflow. This ensures that the developed models and data pipelines run reliably and consistently without manual intervention. The automated execution workflow encompasses the following steps:

1. Creating Databricks Jobs

Convert the finalized notebook into a Databricks job. A job orchestrates the execution of one or more notebooks, automating tasks such as data ingestion, preprocessing, feature engineering, and model training.

2. Defining Workflow Dependencies

Configure the job to include dependencies on other notebooks or tasks as necessary. For example, a job might first execute a data ingestion notebook followed by a model training notebook.

3. Scheduling and Automation

Set up the job to run according to a predefined schedule or trigger based on specific events. This ensures consistent execution of the workflow, whether on a daily basis, in response to new data arrivals, or other criteria.

4. Monitoring and Maintenance

Utilize Databricks’ monitoring tools to track the performance and status of automated jobs, enabling proactive management and troubleshooting.

Automated workflows are supported by Job Clusters in Databricks. Unlike All-Purpose Clusters, Job Clusters are transient and are created specifically for executing jobs. They provide an isolated environment optimized for running automated tasks efficiently, ensuring effective resource utilization and cost management. Job Clusters are ephemeral and cannot be restarted, aligning with their role in supporting transient automated workflows.

Key Features of Databricks Clusters

Databricks Clusters come equipped with several key features that enhance performance, scalability, and security, making them ideal for AI, ML, and generative AI applications.

Cluster Types and Configuration

When creating a cluster, users can choose between a multi-node cluster, where the driver node and the executor nodes reside on separate virtual machines, or a single-node cluster, which combines the driver and executor nodes onto a single machine. This flexibility allows users to tailor cluster configurations based on workload requirements and resource needs.

Auto-Scaling

With auto-scaling, clusters automatically scale up or down based on workload demands. This dynamic scaling optimizes resource utilization and cost by adjusting the number of worker nodes in response to processing needs, ensuring that computational resources are available during peak times and conserved during periods of low demand.

Optimized Apache Spark Environment

Clusters are pre-configured with performance enhancements and the latest versions of Apache Spark, ensuring compatibility and efficiency without the need for manual setup. This optimization includes proprietary improvements that enhance query execution and data processing speeds, enabling faster iteration and more efficient model training.

Security and Isolation

Security and isolation are integral to cluster operations. Clusters run in a secure environment with network isolation and role-based access controls, ensuring that only authorized users can access and manage cluster resources. Users can define custom configurations for clusters, including selecting specific instance types, attaching libraries, and setting environment variables to suit specific workload requirements.

Navigating Databricks Workspaces

A Databricks Workspace is an interactive environment that organizes objects like notebooks, libraries, and experiments into folders and provides tools for collaboration. It serves as the user interface for accessing all Databricks functionalities and managing its resources. The workspace acts as a central hub where users can create and organize their work, access data, and collaborate with team members.

Workspaces enable teams to structure projects logically, maintaining a clear hierarchy of resources that aligns with organizational processes. Centralized management of data access and compliance is facilitated through workspace controls, ensuring that data governance policies are enforced consistently.

The workspace fosters collaboration among data engineers, data scientists, and analysts by providing a shared environment where team members can work together on projects. It enables users to share resources, code, and findings seamlessly. Role-based permissions ensure that users have appropriate access to resources and data, enhancing security and compliance.

Integration with various data sources, libraries, and external tools simplifies workflows and allows teams to leverage existing tools and expertise. The workspace’s organizational structure helps teams manage projects effectively, keeping related notebooks, data, and models grouped logically.

The workspace encompasses several components essential for data professionals, including notebooks which are covered in the next section. The workspace also includes Jobs, which are tools for scheduling and managing automated tasks and workflows. Users can define jobs to run notebooks or scripts at specified intervals or in response to triggers. For data management, the workspace provides tools for browsing and managing datasets, tables, and data sources, simplifying access to the necessary data for analysis.

Collaboration features are built into the workspace, allowing real-time co-editing, commenting, and sharing capabilities. Users can share notebooks and dashboards with colleagues, creating a collaborative environment where insights and knowledge are easily disseminated.

Leveraging Databricks Notebooks

A Databricks Notebook is an interactive, web-based interface that combines code execution, visualizations, and narrative text. Notebooks support multiple languages, such as Python, Scala, SQL, and R, within the same document, allowing users to leverage the strengths of each language as needed. They are essential tools for data exploration, analysis, and machine learning model development, enabling users to iterate quickly and document their processes.

Features of Databricks Notebooks

Databricks notebooks offer a range of features that enhance productivity and collaboration. One of the primary features is multi-language support, which allows users to run code in different programming languages within the same notebook. This flexibility enables data scientists and engineers to utilize the best tools and libraries available for specific tasks.

Real-time collaboration enables multiple users to simultaneously edit and run code, with changes reflected instantaneously. This capability enhances teamwork and accelerates project development by allowing team members to work together closely, regardless of their physical location.

Notebooks provide built-in support for visualizations, enabling users to create charts and graphs to effectively communicate data insights. The ability to include Markdown for formatted text, equations, and images allows users to document methodologies and findings thoroughly within the notebook. This documentation is crucial for knowledge sharing and reproducibility. You can explore the notebooks in this book’s GitHub repository to appreciate the richness of this experience.

Integration with version control systems like Git allows users to track changes, manage branches, and collaborate on code and notebooks. This integration ensures that work is versioned and that team members can work on different features or experiments without conflicts.

Notebooks in ML, AI, and Generative AI

Notebooks are instrumental in the development of machine learning and AI models. They facilitate data exploration by providing an interactive environment where users can examine data distributions, correlations, and anomalies. For feature engineering, notebooks enable the interactive development and testing of data transformations and feature extraction techniques, allowing users to refine their inputs for machine learning models.

During model training and evaluation, notebooks allow for iterative experimentation. Users can train models, adjust hyperparameters, and evaluate performance using various metrics within the same environment. Integration with tools like MLflow enables experiment tracking, where parameters, metrics, and models are logged directly from the notebook, streamlining the machine learning lifecycle.

In the realm of Generative AI, notebooks provide a sandbox for developing and testing generative models such as GPT or variational autoencoders. Users can generate samples from models within the notebook, adjusting parameters and observing outputs in real time. The ability to document the rationale behind model architectures and decisions using Markdown and comments enhances understanding and collaboration among team members.

Integrating Clusters, Workspaces, and Notebooks

The true power of the Databricks platform emerges from the seamless integration of clusters, workspaces, and notebooks. Notebooks are attached to clusters, leveraging their computational resources to execute code efficiently. This integration allows users to run complex data processing and machine learning tasks without worrying about the underlying infrastructure.

Clusters can be dynamically allocated and configured based on the needs of the notebook tasks, optimizing resource usage and ensuring that computational power is available when needed. The workspace provides the organizational structure that houses notebooks and clusters, facilitating collaboration and project management across teams.

This integration supports an end-to-end workflow from data ingestion to model deployment within the Lakehouse environment. Users can perform data processing, analysis, model training, and deployment all within the same platform, reducing operational overhead and simplifying the machine learning pipeline.

Best Practices

To maximize the benefits of Databricks clusters, workspaces, and notebooks, it is important to follow the best practices listed below:

Resource Optimization

Use auto-scaling to adjust resources based on workload demand and terminate idle clusters to manage costs effectively.

Security and Compliance

Implement role-based access controls (RBAC) to manage permissions and monitor permissions within workspaces to protect sensitive data and ensure regulatory compliance.

Version Control and Collaboration

Regularly commit notebooks and code to repositories using version control systems (e.g., Git) to maintain a history of changes, enabling effective collaboration among team members.

Documentation

Use Markdown and inline comments in notebooks to document decisions and methodologies, ensuring proper documentation that facilitates knowledge transfer and reproducibility of experiments.

Code Modularization

Break down code into reusable functions and modules to improve maintainability and readability, encouraging teams to build upon existing work, reduce redundancy, and promote efficiency.

Setting up a Machine Learning Environment

In your journey to harness the full potential of Machine Learning (ML), Artificial Intelligence (AI), and Generative AI within a Lakehouse, establishing a robust and scalable machine learning environment is crucial. Databricks provides a comprehensive platform that simplifies the process of setting up a machine learning environment, enabling data scientists and engineers to focus on model development and innovation rather than infrastructure management. This section guides you through the steps to set up a Databricks Machine Learning environment, highlighting best practices and key considerations.

Understanding the Databricks Machine Learning Workspace

The Databricks Machine Learning workspace is a purpose-built environment within the Databricks platform designed to streamline the end-to-end machine learning lifecycle. It integrates data engineering, exploratory data analysis, model training, deployment, and monitoring into a unified experience. Key features include:

Collaborative Notebooks

Interactive notebooks supporting multiple languages for data exploration and model development.

MLflow Integration

Built-in support for MLflow to manage experiments, track metrics, and deploy models.

Automated Machine Learning (AutoML)

Tools to automatically train and tune models, accelerating the experimentation process.

Feature Store

A centralized repository for storing and sharing machine learning features across projects.

Model Registry

A system for managing model versions, transitions between stages, and deployment workflows.

Prerequisites

Before setting up the Databricks Machine Learning environment, ensure you have:

Access to a Databricks Workspace

You need a Databricks workspace in your cloud environment (AWS, Azure, or GCP).

Permissions

Appropriate permissions to create clusters, manage resources, and access data sources.

Data Sources

Access to the data you plan to use for model training, stored in the Lakehouse.

Creating an Account and a Workspace

This section guides you through the process of creating a Databricks account and workspace. Creating a Databricks account and workspace is the foundational step in leveraging the full potential of the Databricks Lakehouse platform for AI, machine learning, and generative AI development. By carefully selecting your cloud provider, configuring workspace settings, and implementing best practices for security and collaboration, you establish a robust environment that supports efficient data engineering, seamless collaboration, and scalable ML workflows. This setup not only accelerates the development process but also ensures that your AI initiatives are built on a secure, scalable, and highly performant infrastructure.

Creating a Databricks Account

Begin by visiting the Databricks official website¹.

Step 1 - Choose Your Cloud Provider

Databricks is available on multiple cloud platforms, each offering unique integrations and benefits:

Azure Databricks

Integrates with Azure services like Azure Active Directory, Azure Data Lake Storage, and Azure Machine Learning.

AWS Databricks

Integrates with AWS services such as Amazon S3, AWS IAM, and Amazon SageMaker.

Google Cloud Databricks

Connects with Google Cloud services like Google Cloud Storage, BigQuery, and AI Platform.

Select the cloud provider that best aligns with your organization’s infrastructure and strategic goals.

Step 2- Sign up for Databricks

Click on the “Get Started for Free” or “Try Databricks” button on the homepage. You will be prompted to provide essential information (such as your name, company name, work email and phone number) and accept the terms of service to create your account.

After filling in the required details, click “Get Started” or “Sign Up” to proceed.

Step 3 - Verify your Email Address

Databricks will send a verification email to the address you provided. Click on the verification link to activate your account. This step ensures the security and authenticity of your account.

Setting up a Databricks Workspace

Once your account is verified, the next step is to create a Workspace—a collaborative environment where you can manage clusters, notebooks, libraries, and other resources essential for your AI and ML projects.

Step 1 - Access the Databricks Console

After verifying your email, log in to the Databricks Console. If you selected a specific cloud provider during sign-up, you might be redirected to the respective cloud portal (e.g., Azure Portal for Azure Databricks).

Step 2 - Create a New Workspace

In the Databricks Console, locate the option to “Create Workspace” or “New Workspace”. The interface may vary slightly depending on the chosen cloud provider, but the overall process remains consistent.

Step 3 - Configure Workspace Settings

Provide the necessary configurations to set up your workspace:

Workspace Name

Choose a unique and descriptive name for your workspace, reflecting its purpose or the team it serves.

Cloud Region

Select the geographic region closest to your data sources and users to minimize latency and comply with data residency requirements.

Subscription and Billing

Ensure that your workspace is linked to the correct cloud subscription. Review billing details to manage costs effectively.

Network Settings

Depending on your organization’s security policies, you may need to configure Virtual Private Cloud (VPC) settings, network peering, and firewall rules to secure your workspace.

Step 4 - Assign User Permissions

Databricks allows for granular access control to manage who can view, edit, and administer the workspace:

Administrators

Assign workspace administrators who have full control over the workspace settings, cluster configurations, and user permissions.

Users

Add team members who will collaborate within the workspace. Assign appropriate roles (e.g., data scientist, data engineer) based on their responsibilities.

Groups

Utilize cloud provider-specific identity groups (e.g., Azure Active Directory groups) to manage permissions efficiently across multiple users.

Step 5 - Launch the Workspace

After configuring the settings, click “Create Workspace” or “Launch Workspace”. Databricks will initialize the workspace, provisioning the necessary resources and setting up the environment for you and your team to begin their AI, ML and Gen AI projects.

Initial Workspace Configuration and Setup

With your workspace created, it’s essential to configure it to align with your organization’s workflow and security standards. The following steps ensure that your Databricks environment is ready for productive AI and ML development.

Step 1 - Configure Cluster Policies

Databricks offers Cluster Policies to enforce standardized configurations across all clusters, enhancing security and cost management. Navigate to the “Admin Console” and select “Cluster Policies” to create policies that define permissible cluster settings, such as instance types, auto-scaling parameters, and runtime versions.

Step 2 - Setup Repositories

Integrate version control systems like GitHub or GitLab with Databricks to manage your code repositories. This integration facilitates collaboration, version tracking, and continuous integration/continuous deployment (CI/CD) practices.

Step 3 - Establish Data Access and Storage

Connect your workspace to data sources by configuring data storage integrations, such as mounting cloud storage buckets (e.g., Amazon S3, Azure Data Lake Storage, Google Cloud Storage) to the Databricks File System (DBFS).

Step 4 - Enable Security and Compliance Features

Implement security best practices by configuring role-based access controls (RBAC), enabling encryption for data at rest and in transit, and setting up audit logs to monitor user activities. Databricks provides comprehensive security features to help you comply with industry standards and regulatory requirements, such as GDPR and HIPAA.

Step 4 - Create and Configure Notebooks

Begin developing your AI and ML models by creating Databricks notebooks, which serve as interactive environments for writing and executing code. Notebooks support multiple languages, including Python, Scala, R, and SQL.

Setting up Clusters for Machine Learning

Setting up a Databricks cluster tailored for machine learning involves several key steps to ensure optimal performance and compatibility with ML libraries. In this section we will walk through each step.

Accessing your Databricks Workspace

First, you need to log into your Databricks workspace. Log into your cloud platform portal, and navigate to your workspace. Next, click on the Compute icon in the left sidebar of your workspace. You will see any existing clusters that you or other users have created before.

Selecting the Cluster Type

In the top navigation bar, you can choose between a number of cluster types, as is shown in Figure 2-2.

Figure 2-2. - The different cluster types to choose from when setting up a compute type instance.

Here you see the cluster types (All purpose compute and Job compute) we discussed before. SQL warehouses allow you to run SQL query workloads against your data. We will discuss the other options (Vector Search, Pools and Policies) in later chapters. In this case you will be creating an all-purpose cluster, so make sure to leave that tab selected and click the Create Compute button to initiate the cluster creation process.

Multi Node versus Single Cluster

One of the first selections to make when creating a new all-purpose cluster is the multi-node versus single-node radio button, as is shown in Figure 2-3.

Figure 2-3. - Multi Node versus Single node cluster creation.

Multi-Node Clusters consist of separate driver and worker nodes on different virtual machines, making them ideal for large-scale data processing, distributed machine learning model training, and complex analytics workloads. These clusters allow for horizontal scaling by adding more worker nodes, which enhances performance and reduces processing times for data-intensive tasks. Additionally, multi-node clusters offer greater fault tolerance, ensuring higher reliability in production environments. However, they come with increased complexity and higher costs due to the management of multiple instances.

In contrast, Single-Node Clusters integrate both the driver and executor functionalities within a single virtual machine, simplifying cluster management and reducing operational overhead. This configuration is perfect for development, testing, and small-scale data processing tasks such as prototyping machine learning models and performing exploratory data analysis. Single-node clusters are cost-effective and offer fast start-up times, enabling quick iterations and immediate feedback during the development phase. Nevertheless, they are limited in handling large datasets and complex computations, which can lead to performance bottlenecks as project demands grow. Therefore, while single-node clusters are excellent for early-stage development, transitioning to multi-node clusters is recommended for scaling up to production and handling more demanding AI and ML workloads.

Selecting the cluster policy

The first drop-down allows you to select the cluster Policy, as is shown in Figure 2-4.

Figure 2-4. - Policy Options to select from when creating a cluster.

The different Policy options offer several predefined configurations to suit different user needs and workloads. Here is an overview of each option:

Unrestricted

This policy provides maximum flexibility, allowing users to configure compute resources without any limitations. Users can select any node types, cluster sizes, and settings, making it suitable for diverse and dynamic workloads.

Personal Compute

Tailored for individual users, this policy typically provisions a single-node cluster with a larger node type. It’s ideal for personal projects, development, and testing, offering a balance between performance and resource utilization.

Power User Compute

Designed for users requiring more computational power, this policy allows for larger clusters with higher-end node types. It’s suitable for intensive data processing tasks and complex machine learning workloads.

Shared Compute

Intended for collaborative use cases, this policy enables multiple users to share a single cluster. It supports medium to large-scale data processing and is cost-effective for teams working on joint projects.

Legacy Shared Compute

This policy refers to older configurations of shared compute resources. While available for compatibility purposes, it’s recommended to use the current Shared Compute policy for new projects to benefit from the latest features and optimizations.

Selecting the appropriate policy ensures that your compute resources are optimized for your specific requirements, balancing performance, cost, and collaboration needs. For this walkthrough, it is recommended that you select the unrestricted policy to give you the maximum flexibility in your cluster creation.

Selecting the Cluster Access Mode

When configuring a Databricks cluster, selecting the appropriate Access Mode is crucial for aligning the computational resources with your team’s workflow, security requirements and cost considerations. Databricks offers three primary access modes during cluster creation:

Single User

Shared

No Isolation shared

Each access mode caters to different use cases and organizational needs, providing varying levels of resource isolation, security, and collaboration capabilities.

Single User

The Single User access mode dedicates the cluster to a single user. In this mode, only the specified user can attach notebooks and execute code on the cluster. This ensures a fully isolated and secure environment tailored to the individual’s workflow. In single user mode, the supported languages are Python, SQL, Scala and R.

Shared

The Shared access mode allows multiple users to attach to and interact with the same cluster. This mode facilitates collaboration by enabling team members to share computational resources and workspaces seamlessly. In shared mode, the supported languages are Python, SQL and Scala. Note that the ML runtimes (both standard and GPU based) do not support the Shared access mode.

Warning

Shared Access Mode allows multiple users to concurrently attach to and utilize the same cluster, leading to potential resource contention. When multiple users execute heavy ML workloads simultaneously on a shared cluster, it can result in diminished performance, increased latency, and inefficient utilization of computational resources. By restricting ML Runtimes to non-shared access modes, Databricks ensures that the clusters are dedicated to single-user or isolated workflows, thereby maintaining high performance and stability for ML operations.

No Isolation Shared

The No Isolation Shared access mode offers a hybrid approach where multiple users can share the same cluster without stringent isolation. This mode is designed for scenarios where collaboration is essential, but full isolation is not a primary concern. This mode supports the Python, SQL, Scala and R programming languages.

Access Mode Selection Advice

Choosing the right access mode depends on several factors, including:

Collaboration Needs

Determine whether the workload requires individual focus or team collaboration.

Security Requirements

Assess the sensitivity of the data and the need for isolation to protect against unauthorized access.

Cost Constraints

Balance the need for dedicated resources against budgetary considerations to optimize cost efficiency.

Performance Expectations

Consider the computational demands of your workloads and how resource sharing might impact performance.

Recommendation

Use Single User mode when security and resource isolation are paramount, such as handling sensitive data or performing critical computations. Single User mode is also recommended for any ML development work with the ML runtimes.

Opt for Shared mode to facilitate collaboration among team members working on joint projects, ensuring efficient resource utilization while maintaining a degree of isolation.

Choose No Isolation Shared mode for flexible, low-security environments where ease of access and cost savings take precedence over stringent security and performance guarantees. In a real-life scenario, it is hard to think of any reason to select this option. This mode is not supported by Unity Catalog which is another reason to forgo this option.

If you will be the only one using the cluster, you should select Single User mode, if you are going to work in a team on non ML development, you should select Shared Mode. For any ML work, you should choose Single User Mode.

Runtime Selection

The Machine Learning (ML) Runtimes are specially optimized versions of the Databricks Runtime, equipped with a comprehensive suite of ML libraries and frameworks such as TensorFlow, PyTorch, Scikit-learn, and MLflow. These runtimes are designed to facilitate efficient development, training, and deployment of machine learning models by providing enhanced performance, specialized tools, and seamless integration with ML workflows.

When selecting a runtime, you have a choice to select between Standard and ML, as is shown in Figure 2-5.

Figure 2-5. - Create cluster runtime selection

Databricks Runtime (Standard)

The Databricks Runtime (Standard) is the foundational runtime environment optimized for a wide range of data engineering and analytics tasks. It incorporates the latest stable release of Apache Spark, along with performance enhancements and essential libraries for data processing, transformation, and visualization.

The standard runtime’s key features include:

Optimized Spark Engine

Enhances the performance of Spark jobs through proprietary optimizations, reducing execution time and resource consumption.

Pre-installed Libraries

Includes essential libraries such as Delta Lake for reliable data storage, and integrations with BI tools for seamless data visualization.

Compatibility

Supports multiple programming languages, including Python, Scala, R, and SQL, providing flexibility for diverse development needs.

The main use cases for the standard runtime are:

Data ingestion and transformation

ETL (Extract, Transform, Load) pipelines

Interactive data analysis and visualization

Long-Term Support (LTS) versus Non-LTS Runtimes

Databricks offers two primary categories of standard runtimes: Long-Term Support (LTS) Runtimes and Non-LTS Runtimes. Understanding the distinctions between these runtimes and their respective use cases is essential for optimizing your development workflows and maintaining robust, scalable solutions.

LTS Runtimes are designed to provide a stable and reliable foundation for production environments over an extended period. These runtimes receive regular updates that focus primarily on security patches, critical bug fixes, and essential performance enhancements, without introducing significant changes to the core functionalities or libraries. The primary advantage of LTS Runtimes is their consistency and reliability, making them ideal for organizations that prioritize stability and need to minimize disruptions caused by frequent runtime changes.

Non-LTS Runtimes, often referred to as latest or standard runtimes, incorporate the newest features, library updates, and performance improvements introduced by Databricks. These runtimes are ideal for experimentation, development, and leveraging the latest advancements in AI and ML technologies. While they offer enhanced capabilities and access to the most recent tools, Non-LTS Runtimes may also introduce changes that could affect compatibility or require adjustments in your workflows.

Databricks Runtime for Machine Learning

The Databricks Runtime for Machine Learning is specifically engineered to accelerate ML workflows by bundling advanced ML libraries and frameworks. This runtime is ideal for data scientists and ML engineers who require a comprehensive suite of tools to develop, train, and deploy machine learning models efficiently.

The key features of the machine learning runtime are:

Integrated ML Libraries

Comes with pre-installed libraries such as TensorFlow, PyTorch, Scikit-learn, XGBoost, and MLflow, facilitating end-to-end ML workflows.

Optimized Environments

Configured to support distributed training and hyperparameter tuning, leveraging Databricks’ scalable infrastructure.

Seamless MLflow Integration

Enhances model tracking, experiment management, and deployment processes, ensuring reproducibility and streamlined collaboration.

The main use cases for the ML runtime are:

Developing and training machine learning models

Experimenting with different algorithms and hyperparameters

Managing the ML lifecycle from experimentation to deployment

GPU Support

The observant reader has likely noticed that in figure 254, for every runtime version you have a choice between the regular ML runtime and a GPU runtime. For workloads that demand intensive computational power, such as deep learning and generative AI, the Databricks Runtime with GPU Support provides the necessary infrastructure to handle parallel processing tasks efficiently. This runtime leverages Graphics Processing Units (GPUs) to accelerate data processing and model training, significantly reducing computation times.

The key features of this runtime are:

GPU-Enabled Instances

Utilizes high-performance GPU instances (e.g., NVIDIA Tesla V100, A100) to expedite computationally heavy tasks.

Deep Learning Frameworks

Optimized for frameworks like TensorFlow and PyTorch, enabling faster training of deep neural networks.

Enhanced Parallelism

Supports distributed GPU training, allowing simultaneous processing across multiple GPUs for large-scale models.

The use cases targeted by the GPU runtimes are:

Training deep learning models for computer vision, natural language processing (NLP), and generative tasks

Performing real-time data processing and inference

Running simulations and complex algorithmic computations

Photon Acceleration

In the evolving landscape of AI, machine learning (ML), and generative AI, performance and efficiency are paramount. Photon Acceleration is a cutting-edge feature within the Databricks Runtime that significantly enhances the speed and efficiency of data processing and query execution. When creating a Databricks cluster, enabling Photon Acceleration via the "Photon Acceleration" checkbox unlocks these performance benefits, making it an invaluable tool for data scientists and engineers working on demanding AI and ML workloads.

Photon is Databricks’ native vectorized execution engine, designed to deliver high-performance query processing by leveraging modern CPU architectures. Built into the Databricks Runtime, Photon transforms how data is processed within the Lakehouse environment by optimizing Spark’s execution framework. It achieves this through advanced techniques such as Single Instruction, Multiple Data (SIMD) parallelism and efficient memory management, which collectively reduce query latency and increase throughput.

The key benefits of using Photon acceleration are:

Enhanced Query Performance

Photon drastically reduces query execution times by optimizing data processing pathways. This acceleration is crucial for iterative ML model training and real-time data analytics, where rapid data manipulation is essential.

Improved Resource Utilization

By efficiently managing CPU resources and memory, Photon ensures that computational power is maximized, reducing idle times and improving overall cluster efficiency. This leads to cost savings by minimizing the need for additional computational resources.

Scalability and Throughput

Photon’s vectorized execution engine allows clusters to handle larger datasets and more complex queries without a proportional increase in resource consumption. This scalability is vital for generative AI tasks that involve processing extensive data inputs and generating sophisticated outputs.

Seamless Integration

Photon integrates seamlessly with existing Spark-based workflows, enabling users to benefit from performance enhancements without altering their codebases. This ease of adoption ensures that teams can quickly leverage Photon’s capabilities within their established processes.

Driver and Worker Nodes

Understanding the distinct roles of Driver and Worker Nodes is fundamental to configuring a Databricks cluster effectively.

The Driver Node serves as the cluster’s orchestrator, managing the SparkContext, scheduling tasks, and coordinating job execution across Worker Nodes. Its responsibilities include maintaining the Spark application’s state, distributing tasks, and handling communications between the cluster and external systems.

In contrast, Worker Nodes are tasked with executing the data processing and computations assigned by the Driver Node. Their duties encompass performing transformations and actions on data, executing user-defined functions and machine learning algorithms, and managing memory and storage for intermediate data.

Configuring Driver and Worker Nodes

When creating a Databricks cluster, users must configure both Driver and Worker Nodes to align with their specific workload requirements. Let’s now discuss the key considerations and options available for each node type.

Selecting Instance Types

Choosing the right instance types for Driver and Worker Nodes is paramount for ensuring that the cluster can handle the computational demands of AI, ML, and generative AI tasks efficiently. You typically have a choice between four different instance families:

General-Purpose Instances

Balanced compute, memory, and storage, suitable for a wide range of workloads. Examples: n1-standard-8 (8 vCPUs, 30 GB RAM)

Compute-Optimized Instances

High-performance CPUs for CPU-intensive tasks. Examples: c2-standard-16 (16 vCPUs, 60 GB RAM)

Memory-Optimized Instances

Enhanced memory capacity for memory-intensive applications. Examples: m1-ultramem-40 (40 vCPUs, 961 GB RAM)

GPU-Enabled Instances

Equipped with Graphics Processing Units (GPUs) for parallel processing tasks essential in deep learning and generative AI. Examples: a2-highgpu-8g (96 vCPUs, 680 GB RAM, 8 NVIDIA A100 GPUs)

Driver Node Configuration

The Driver Node should typically mirror the Worker Nodes in terms of instance type to prevent bottlenecks. For instance, if using GPU-enabled Worker Nodes, the Driver should also be GPU-enabled. Ensure the Driver Node has sufficient memory and CPU resources to manage the cluster efficiently, especially for large-scale ML tasks.

Worker Node Configuration

Select instance types that can scale based on workload demands. For AI and ML workloads, memory-optimized or GPU-enabled instances are often preferred. Depending on the task, users might opt for a mix of instance types within the same cluster to balance cost and performance.

Configuring Cluster Size

Determining the optimal number of Worker Nodes is essential for balancing performance and cost within a Databricks cluster. To accommodate varying computational demands, Databricks offers flexible scaling options that allow clusters to adapt dynamically to workload fluctuations. These scaling strategies ensure that resources are utilized efficiently, enhancing both the performance of data processing tasks and the cost-effectiveness of cluster operations.

One of the primary scaling options is Auto-Scaling, which automatically adjusts the number of Worker Nodes based on the current workload. This offers significant benefits, including cost efficiency and performance optimization. During periods of low demand, Auto-Scaling reduces the number of Worker Nodes, thereby lowering operational costs by minimizing resource usage. Conversely, when workloads increase, such as during peak data processing or intensive machine learning model training, Auto-Scaling increases the number of Worker Nodes to handle the additional load effectively, ensuring that jobs are executed in a timely manner.

Configuring Auto-Scaling involves setting the minimum and maximum number of Worker Nodes, which establishes the boundaries for scaling behavior. Additionally, scaling triggers can be defined based on specific metrics such as CPU utilization, memory usage, or job queue length. These triggers enable the cluster to respond intelligently to real-time resource demands, maintaining optimal performance without manual intervention.

Alternatively, Databricks provides the option of maintaining a Fixed Cluster Size, where the number of Worker Nodes remains constant regardless of workload variations. This approach is particularly beneficial for workloads that are stable and predictable, ensuring that a consistent level of computational resources is always available. Fixed Cluster Size offers predictable performance, as the cluster’s resource allocation does not fluctuate with varying demands, which is crucial for long-running tasks that require uninterrupted processing power.

Additionally, this mode simplifies cluster management by reducing the complexity associated with dynamic scaling, making it easier to oversee and maintain the cluster. Fixed clusters are ideal for scenarios where workloads do not experience significant fluctuations, such as ongoing data ingestion processes or continuous model training pipelines that operate at a steady pace. By choosing a Fixed Cluster Size, organizations can ensure reliability and consistency in their computational environment, aligning resource allocation with their specific operational requirements.

Automatic Termination After Inactivity

To prevent clusters from running idle and incurring unnecessary costs, Databricks provides an automatic termination feature. During cluster setup, you can define an inactivity period (in minutes). If the cluster remains inactive, without any executed commands or jobs beyond this specified duration, Databricks will automatically terminate the cluster

Final Configuration

The final configuration screen for a typical machine learning development cluster is shown in Figure 2-6. Notice the summary display on the right which provides a quick summary of your configuration, a quick overview of your cluster, whether or not Unity Catalog is supported by your cluster configuration, and the total number of Databricks Units (DBUs) per hour.

A DBU is a standardized measure of processing power on the Databricks platform. It represents the amount of computational resources required to execute a specific workload, encompassing factors like CPU, memory, and I/O operations. By abstracting these underlying resources into DBUs, Databricks provides a consistent and straightforward metric for users to gauge and manage their computational needs.

Databricks employs a consumption-based pricing approach, charging users based on the number of DBUs utilized. This model ensures that you pay only for the processing power you consume, aligning costs directly with usage. The pricing per DBU can vary depending on factors like the cloud provider (e.g., AWS, Azure, Google Cloud) and the specific instance types selected.

Figure 2-6. Final Cluster Configuration

Extending the ML Runtime with Additional Libraries

While the Databricks Machine Learning (ML) Runtime includes essential libraries like TensorFlow, PyTorch, Scikit-learn, XGBoost, and MLflow, certain AI and ML projects require additional libraries. The pre-installed suite offers a strong foundation for diverse tasks, enabling efficient model development, training, and deployment. However, the rapidly evolving nature of AI and ML often necessitates specialized tools, custom functionalities, or the latest advancements not covered by the default runtime.

Specialized and Cutting-Edge Libraries

AI and ML fields continuously introduce new libraries and frameworks to tackle specific challenges and enhance capabilities. Projects involving advanced natural language processing (NLP), computer vision, or reinforcement learning can benefit from libraries beyond the standard ML Runtime. For instance, integrating Hugging Face Transformers for sophisticated NLP tasks, Detectron2 for object detection, or Ray for distributed computing can provide tailored functionalities essential for unique project requirements. Installing these specialized libraries ensures that the development environment remains adaptable and leverages the latest AI and ML innovations.

Custom and Proprietary Libraries

Many organizations develop proprietary libraries or custom extensions to meet specific business needs or implement unique algorithms. These custom libraries may include specialized data processing techniques, unique model architectures, or custom evaluation metrics critical to organizational operations. By adding these libraries, teams can seamlessly incorporate proprietary functionalities into their Databricks environment, ensuring powerful and standardized AI and ML workflows. Additionally, custom libraries promote better code reuse, maintainability, and collaboration across various projects and teams within the organization.

Summary

The Databricks Intelligence Platform provides a robust, scalable, and secure environment tailored for advanced AI, machine learning and generative AI development. By unifying data engineering, data science, and ML workflows within a single Lakehouse architecture, Databricks eliminates data silos and fosters seamless collaboration across teams. The platform’s optimized runtime environments, versatile cluster configurations, and comprehensive machine learning features empower organizations to efficiently develop, train, and deploy sophisticated models. Additionally, Databricks’ security and governance features ensure that sensitive data remains protected, while its collaborative tools enhance productivity and innovation. Together, these elements make Databricks an ideal foundation for organizations aiming to leverage large-scale data for innovative AI-driven initiatives, driving both technological advancement and business value.

¹ Databricks Web Site: https://www.databricks.com/

Chapter 3. An Intro to ML on Databricks

A Note for Early Release Readers

With Early Release ebooks, you get books in their earliest form—the author’s raw and unedited content as they write—so you can take advantage of these technologies long before the official release of these titles.

This will be the 3rd chapter of the final book. Please note that the GitHub repo will be made active later on.

If you have comments about how we might improve the content and/or examples in this book, or if you notice missing material within this chapter, please reach out to the editor at jbleiel@oreilly.com.

In this chapter, you will explore the essential components of machine learning workflows tailored for the Databricks platform. You’ll be guided through an end-to-end process, starting with exploratory data analysis (EDA) to build an understanding of your dataset, followed by feature engineering to transform raw data into a feature matrix for building robust machine learning models.

Next, you’ll focus on model development, including creation and training. After training, you’ll learn about model deployment and serving, setting up the necessary infrastructure for real-time inference to provide actionable insights at scale.

We will then cover MLflow, an open-source platform, which streamlines the machine learning lifecycle from experimentation to deployment. Throughout this book, you’ll see how MLflow simplifies the ML workflow within Databricks, empowering you to efficiently build, track, and serve machine learning models.

By the end of this chapter, you’ll have a comprehensive understanding of the entire ML lifecycle in a Databricks Lakehouse environment, from data preparation to production-grade model serving.

The End-to-End Machine Learning Environment

In today’s data-driven landscape, the ability to efficiently build, deploy, and manage machine learning models is crucial for businesses aiming to stay competitive. Databricks offers a unified platform that streamlines the entire machine learning workflow, from data preparation to model deployment, ensuring scalability, collaboration, and performance throughout the process. The end-to-end workflow is shown in Figure 3-1.

Figure 3-1. The end-to-end Machine Learning Workflow

This section delves into how Databricks facilitates each stage of the ML lifecycle, ensuring scalability, collaboration, and performance.

Data Preparation

High-quality data is the foundation of any successful machine learning project. Databricks integrates seamlessly with a wide variety of data sources, enabling efficient data ingestion and preparation. Whether your data sources reside in cloud storage like AWS S3, Azure Blob Storage, Google Cloud Storage; relational databases such as MySQL, PostgreSQL, SQL Server; or streaming platforms like Apache Kafka, Amazon Kinesis, and Azure Event Hubs, Databricks provides connectors to ingest and process data efficiently. It also supports diverse file formats including CSV, JSON, Parquet, and Avro, ensuring flexibility in handling different data types.

By offering comprehensive connectivity options, Databricks simplifies the data ingestion process, allowing data scientists and engineers to focus on data analysis and model development rather than the complexities of data retrieval. This streamlined approach accelerates the development of machine learning models and unlocks the full potential of your data within the lakehouse environment.

Delta tables and the Feature Store enhance consistency, scalability, and efficiency. Delta tables ensure transactional data handling, supporting incremental updates and data versioning, making it easier to maintain clean, up-to-date datasets. The Feature Store centralizes feature engineering, ensuring consistent transformations and enabling efficient feature reuse across different datasets. Together, they streamline the entire data preparation process, facilitating robust and reliable data engineering pipelines.

Delta Tables

Once data is ingested, storing it in a format that supports both analytical and transactional workloads is essential. Databricks uses Delta Lake’s Delta Tables to bring ACID transactions and schema enforcement to big data workloads, ensuring data integrity. ACID compliance guarantees that concurrent reads and writes occur smoothly without data corruption, which is crucial for maintaining accurate data in complex machine learning workflows.

Delta Tables also enhance performance by optimizing queries with built-in indexing and caching mechanisms, allowing faster data retrieval and analytics. This reduces query latency and improves overall system efficiency, especially when working with large datasets. Delta Lake supports features like time travel for accessing previous data versions for auditing or reproducing experiments, and schema evolution to modify data schemas without compromising integrity.

By using Delta Tables, you maintain clean, versioned data ready for machine learning tasks. Their scalability, reliability, and performance ensure your data pipeline is robust and efficient, allowing your team to focus on model development and analysis rather than data management challenges.

Feature Store

Effective management of features, the variables that machine learning models use to make predictions, is crucial for ensuring model accuracy and consistency. The Databricks Feature Store provides a centralized repository for managing machine learning features, allowing teams to register, discover, and serve features consistently across various stages of model development and deployment.

Registering features in the Feature Store enables data scientists and engineers to store and document them for reuse, promoting collaboration and reducing redundancy. Team members can search for and utilize existing features, avoiding duplication of effort and fostering a more efficient workflow. Consistent serving of features during both training and inference ensures that the same feature definitions and values are used throughout the machine learning pipeline, eliminating discrepancies that might arise from using different versions of a feature.

The Databricks Feature Store supports feature versioning and integrates seamlessly with other components of the Databricks ecosystem, such as Delta Lake for data storage and management, ensuring that feature data is reliable and performant.

Model Training

Databricks provides robust tools for model training. This facilitates efficient workflows for model development and data analysis.

Databricks Notebooks offer an interactive and collaborative environment for developing and visualizing machine learning models using languages like Python, R, and Scala. Notebooks seamlessly integrate code execution, rich text, and visualizations, supporting efficient workflows for both model development and data analysis within the Databricks ecosystem. Integration with version control systems like Git enables users to track changes, manage different code versions, and ensure that all modifications are documented and reversible. The notebooks also support data visualization through built-in libraries, which is crucial for exploratory data analysis and informed decision-making in model development.

For those looking to expedite model development, Databricks AutoML provides automated machine learning capabilities that simplify and accelerate the creation of machine learning models. It automates feature engineering, model selection, and hyperparameter tuning, evaluating multiple algorithms and configurations to find the best-performing model. This tool lowers the barrier to entry for those new to machine learning while providing advanced tools that benefit experienced practitioners, ultimately accelerating the model development lifecycle.

MLflow, integrated into Databricks, is designed to manage the end-to-end machine learning lifecycle. We will cover MLflow in detail in the next chapter.

Model Creation and Validation

After experimenting and identifying the best-performing model, the next step is to prepare it for deployment. This involves model serialization—saving the trained model in a format suitable for deployment, such as the MLflow Model format—and packaging all necessary dependencies to prevent issues related to missing packages or version mismatches in the deployment environment. Validation is also a critical component, involving testing the model with validation datasets to confirm it performs as expected. Databricks integrates with MLflow Models to streamline this entire process, enabling a seamless transition from development to deployment.

Model Productionalization

Deploying machine learning models into production environments is the final step in the end-to-end workflow, and Databricks supports this process through various methods:

Batch Inference is ideal for applications where real-time predictions are not required, batch inference allows for periodic predictions on large volumes of data. Databricks supports batch inference by enabling users to run prediction jobs at regular intervals using Databricks Jobs. These scheduled jobs leverage scalable, distributed computing to handle large datasets efficiently, integrating with Delta Lake for seamless reading and writing of data.

For applications requiring real-time predictions, streaming inference is essential. Databricks enables streaming inference through Structured Streaming, allowing users to process data streams with high-level APIs. This ensures low latency and fault tolerance, critical for time-sensitive applications like fraud detection or real-time recommendations.

Databricks offers online serving endpoints to make models accessible for real-time scoring via APIs. These endpoints expose models as RESTful APIs that can be called directly from applications, providing immediate predictions. The platform ensures scalability by automatically adjusting resources based on demand and includes monitoring features to track performance metrics and logs. Security measures are also implemented to protect access to the models.

Whether through batch processing, streaming data, or real-time API calls, Databricks ensures that models can be operationalized efficiently, allowing businesses to quickly derive value from their data in the lakehouse environment.

Managing the ML Lifecycle with MLflow

In modern machine learning workflows, managing experiments, ensuring reproducibility, and efficiently deploying models are critical challenges. MLflow is an open-source platform that addresses these challenges by providing comprehensive solutions for the entire machine learning lifecycle. Integrated seamlessly into Databricks, MLflow offers tools for tracking experiments, packaging code into reproducible runs, and sharing and deploying models, thus streamlining the path from experimentation to production.

The model development lifecycle with MLflow is shown in Figure 3-2.

Figure 3-2. The end-to-end development workflow with MFflow

MLflow Components

MLflow comprises several key components that work together to facilitate an efficient machine learning workflow:

Model Tracking

This component logs and queries experiments, tracking parameters, metrics, and artifacts associated with models.

Model Management

Models are managed in a standardized, serialized format suitable for prediction, allowing for easy deployment across various platforms.

Model Serving

This component makes models available for batch and streaming inference, facilitating integration into production environments.

MLflow Projects

Mlflow Projects package code and dependencies to recreate the complete environment needed for training, ensuring consistency across different setups.

Model Registry

The Model Registry enables you to register models to manage their lifecycle and deployment stages, providing a centralized repository for model management.

Model Tracking

Model tracking is fundamental in MLflow, providing a centralized system to log parameters, such as hyperparameters and configurations. It also provides metrics like accuracy and loss, and artifacts including model binaries and plots, and source code versions. By recording each run, MLflow enables reproducibility, comparison, and collaboration. This detailed logging makes it easier to replicate results, understand why certain models perform better, and share findings with team members.

An experiment in MLflow serves as a container for runs, which are individual executions of model code with specific parameters. Experiments help organize runs and control access, facilitating visualization, search, and comparison in the MLflow user interface. MLflow supports two types of experiments: Workspace Experiments and Notebook Experiments. Workspace Experiments are explicitly created and accessible across notebooks and code within the Databricks workspace, making them ideal for collaborative projects. Notebook Experiments are implicitly created when running code in a notebook without specifying an experiment, suitable for individual exploration.

A run captures execution details, including start and end times, parameters, metrics, artifacts, tags, and source information, aiding in reproducibility and analysis. MLflow provides APIs in several programming languages, Python, REST, R, and Java, allowing flexible integration into various development environments.

Model Management

MLflow simplifies model management and deployment by providing a consistent approach to packaging models, regardless of the machine learning library used. It supports models from popular libraries such as scikit-learn, XGBoost, PyTorch, and TensorFlow. When you log a model with MLflow, it is serialized into a standardized format that includes:

The model binary, which is the actual model file.

A Conda environment specification that lists all dependencies required to run the model.

An MLmodel file containing metadata about the model’s format and other relevant information.

This standardized packaging ensures that models can be easily loaded and served in different environments without compatibility issues.

Model Serving

MLflow supports both batch and real-time inferencing, catering to different application needs. Batch inferencing is suitable for generating predictions offline on large datasets. MLflow allows you to load serialized models into batch processing jobs using the MLflow API. This leverages distributed computing frameworks like Apache Spark to process large volumes of data efficiently.

For applications requiring immediate predictions, real-time inferencing is essential. MLflow offers two approaches:

Classic Model Serving

This method deploys the model to a REST API endpoint on a single-node cluster managed by MLflow. It is suitable for low to moderate load scenarios where you need more control over the serving environment.

Serverless Real-Time Inferencing

Currently in public preview, this feature provides scalable REST API endpoints without the need to manage servers or clusters. The model is hosted on scalable infrastructure managed by Databricks, with resources that automatically adjust based on incoming traffic. This approach is ideal for production environments requiring scalability and reliability, enhancing cost efficiency by only charging for compute resources used during model serving.

MLflow Projects

MLflow Projects standardize the packaging of machine learning code and its environment into reusable and reproducible forms. A project typically includes:

An MLproject file, which is a YAML file defining the project structure, entry points, parameters, and environment specifications.

A Conda environment file specifying the dependencies needed to run the project.

The source code, such as scripts or notebooks containing the model training code.

Any necessary data files, often referenced rather than included directly.

By using MLflow Projects, you ensure reproducibility, allowing anyone to recreate the exact environment and reproduce your results. This facilitates collaboration, as you can share your project with other data scientists who can run experiments with different parameters or data while maintaining consistency. It also simplifies deployment by packaging the project for production environments, ensuring all dependencies are met.

Projects can be run using the MLflow command-line interface (CLI) or programmatically via the MLflow API. When you run a project, MLflow handles environment setup by creating a new environment based on the specified dependencies. It accepts parameters and passes them to the entry point script, and if you’re using MLflow tracking APIs within the code, it automatically tracks the run.

Model Registry

The MLflow Model Registry is a centralized repository and collaborative hub where you can register, manage, and track the lifecycle of machine learning models. It provides a structured way to:

Add Trained Models

After training and evaluation, models can be registered in the Model Registry as new versions.

Enforce Model Version Control

Keep track of different versions of a model, tracking changes over time and ensuring traceability.

Manage Lifecycle Stages

Assign models to stages such as “Staging,” “Production,” or “Archived,” controlling the progression of models through testing and deployment.

Annotate Models

Add descriptions, comments, and metadata to models for better context and understanding.

Establish Access Restrictions

Implement permissions to manage who can register, update, or deploy models, enhancing governance and security.

By managing the model lifecycle through the Model Registry, you ensure that only thoroughly tested and approved models reach production. Models in the “Staging” stage undergo testing and validation, while those promoted to “Production” are approved for deployment in production environments. Models that are no longer in active use can be “Archived” but retained for record-keeping or potential future use.

Centralizing model management supports team collaboration, allowing multiple team members to work on models, track changes, and share updates. It enhances auditability by maintaining detailed records of model versions and changes, which is crucial for compliance purposes. The registry also promotes reusability by making models discoverable and accessible within the organization.

Example Workflow

Consider a typical ML workflow:

1. Training and Logging

A data scientist trains a new model and logs it using MLflow, capturing all relevant parameters and metrics, and saving them as model tracking artifacts.

2. Packaging

The best performing model is packaged as an MLflow project, ensuring that we can easily reconstitute the model, and repeat our experiments.

3. Registration

The trained model is registered in the Model Registry as a new version under a specific model name.

4. Staging

The model is moved to the “Staging” stage for testing and validation, where it undergoes thorough testing by the QA team or automated tests.

5. Production Deployment

Upon successful validation, the model is promoted to “Production” and deployed to production environments.

6. Monitoring and Iteration

The model’s performance is monitored in production. If issues arise or improvements are identified, a new model version is trained, and the cycle repeats.

By following this structured approach, organizations can ensure that only high-quality models make it to production, maintaining reliability and performance in their machine learning applications.

Summary

In this chapter, you explored MLflow, a vital platform for managing the entire machine learning lifecycle, and how it integrates seamlessly with the Databricks Lakehouse environment. You began by walking through the key components of a machine learning workflow, starting with exploratory data analysis (EDA) and feature engineering, where you transformed raw data into meaningful features ready for model development.

You then delved into model creation, training, and deployment, learning how Databricks simplifies these stages through its unified platform. Tools like Delta Tables and the Feature Store were introduced to help you ensure data integrity, scalability, and consistency during data preparation. You also saw how MLflow streamlines your machine learning process, enabling you to efficiently track experiments, package reproducible runs, and manage models in a centralized location.

You explored the core components of MLflow, including Model Tracking, Model Management, Model Registry, Model Serving, and MLflow Projects, and learned how these tools address common challenges like experiment organization, model versioning, and deployment for batch or real-time inference. You followed practical workflows, such as training models, validating them in staging environments, and deploying them to production systems.

By the end of this chapter, you gained a comprehensive understanding of how Databricks and MLflow empower you to accelerate your machine learning workflows, collaborate effectively, and deploy production-grade models efficiently. This knowledge prepares you to dive deeper into MLflow’s advanced capabilities in the chapters ahead.

In the next chapter, you will put this knowledge into action by building an end-to-end example of the machine learning workflow. From data preparation to model deployment, you’ll gain hands-on experience implementing the concepts introduced in this chapter, providing you with a complete understanding of how these tools come together in a real-world scenario.

Chapter 4. Machine Learning Use Case with MLflow

A Note for Early Release Readers

With Early Release ebooks, you get books in their earliest form—the author’s raw and unedited content as they write—so you can take advantage of these technologies long before the official release of these titles.

This will be the 4th chapter of the final book. Please note that the GitHub repo will be made active later on.

If you have comments about how we might improve the content and/or examples in this book, or if you notice missing material within this chapter, please reach out to the editor at jbleiel@oreilly.com.

In this chapter, we explore a practical machine learning use case using MLflow on the Databricks platform. The use case involves predicting hotel booking cancellations, leveraging the Kaggle hotel_bookings.csv dataset. By building a model that accurately forecasts cancellations, hotels can optimize bookings, reduce no-shows, and improve occupancy management.

This hands-on example illustrates the entire ML lifecycle in a lakehouse environment, including data exploration, feature engineering, model training, and deployment. You’ll also learn how to use MLflow for experiment tracking, model versioning, and deployment, enabling efficient and reproducible machine learning workflows on Databricks.

By the end of this chapter, you’ll have a clear understanding of how MLflow and Databricks combine to streamline machine learning development and operationalization, driving faster insights and better business outcomes.

Use Case Overview

Your journey begins with Exploratory Data Analysis (EDA), where you will examine and visualize the data to gain a comprehensive understanding of its structure and characteristics. By conducting EDA, you will uncover patterns, anomalies, and relationships within the dataset, laying a solid foundation for the subsequent stages.

Once you have a deep understanding of the data, you will move on to the Feature Engineering phase, where you will transform raw data into a feature matrix optimized for machine learning models. This step involves creating new features, encoding categorical variables, and scaling numerical variables, all of which prepare the dataset for model building.

With the dataset ready, you will set up an MLflow experiment to build, train and evaluate multiple machine learning models. Using MLflow’s Autologging feature, you will automatically capture key metrics, parameters, and artifacts for each model run, making it easy to track and compare performance. By logging all experimental details, you will seamlessly iterate on model training while maintaining an organized history of each model’s results.

After identifying the best-performing model, you will register the model in the MLflow Model Registry, ensuring it is versioned, documented, and available for collaborative review. The registered model can then be deployed as a REST API endpoint to enable real-time inferencing, allowing you to integrate your model’s predictions into downstream applications and workflows.

This end-to-end pipeline, from data exploration to model deployment, demonstrates the power and flexibility of a lakehouse architecture with MLflow. The pipelines enables a seamless transition from raw data to operationalized machine learning models, empowering you to create, deploy, and monitor advanced AI solutions with agility and scale.

This workflow is captured in Figure 4-1.

Figure 4-1. End-to-End Machine Learning Use Case

You will work through all the steps of the model workflow in the next sections of this chapter.

Exploratory Data Analysis

Building an effective machine learning model begins with thorough data preparation. In this section, you will load the dataset that serves as the foundation for your machine learning model within the Lakehouse environment. Leveraging the interactive capabilities of Databricks notebooks, you will conduct an Exploratory Data Analysis (EDA) to understand the dataset’s structure, identify key features, and detect any anomalies or missing values that could influence model performance. This analysis will provide the necessary insights to inform the subsequent data preprocessing steps.

Once you have a firm grasp of the dataset, you will move on to data preprocessing, which includes cleaning, transforming, and encoding the raw data into a format optimized for training your machine learning model.

Loading the Dataset

To start your machine learning project, load the dataset into the Databricks File System (DBFS), the foundational storage layer in Databricks. Use the DBFS Browser, a built-in graphical interface, to easily upload and manage files without needing code or command-line utilities. This intuitive tool allows you to focus on data exploration and analysis rather than the technicalities of data loading.

Enabling the DBFS Browser

In recent Databricks versions, the DBFS Browser is disabled by default for security reasons. To enable it, go to Settings via the user icon in the top-right corner, click Advanced under Workspace Admin, and toggle the DBFS File Browser setting.

Once enabled, you can easily manage datasets visually, a helpful step when starting your machine learning workflow. Next, download the hotel_bookings.csv file from Kaggle and upload it to Databricks.

Uploading the Dataset to DBFS

To upload the hotel_bookings.csv file into the /FileStore/datasets using the DBFS Browser:

Download the file from Kaggle to your local machine.

In Databricks go to Catalog in the left navigation pane and click Browse DBFS.

Under the FileStore directory, create a subdirectory named datasets.

Navigate to /FileStore/datasets,click Upload and select the hotel_bookings.csv file (see Figure 4-2)

Figure 4-2. Uploading the hotel-booking.csv file.

Once uploaded, you’re ready to begin Exploratory Data Analysis (EDA) to examine the dataset, visualize trends, and identify any issues for preprocessing and model development.

Exploring the Dataset

The Exploratory Data Analysis (EDA) notebook is available at /chapter04/04-01-Exploratory Data Analysis in the GitHub repository. While this book focuses on machine learning techniques, running the EDA notebook is highly recommended as it lists key questions about the data along with corresponding answers, while also validating or invalidating a set of hypotheses to guide feature engineering and model building.

Feature Engineering

The Feature Engineering notebook, available at /chapter04/04-02-Feature Engineering in the book’s Git repository, focuses on transforming raw data into meaningful features to enhance model performance. It covers essential tasks like handling missing values, removing duplicates, encoding categorical variables, and reducing redundancy through cross-correlation analysis. The complete table of contents for the feature engineering notebook is shown in Figure 4-4.

Figure 4-3. Feature Engineering steps

By following the notebook, you’ll prepare the hotel_bookings.csv dataset for machine learning and save the processed DataFrame in Delta format at dbfs:/FileStore/datasets/hotel_bookings.delta. Reviewing this notebook is highly recommended to understand the steps crucial for building effective models.

Model Training

Before training any machine learning models, you’ll first set up an MLflow experiment to ensure that all your model runs are tracked and organized. This experiment will serve as the container for capturing essential details such as metrics, parameters, and model artifacts throughout the modeling process.

Creating an MLflow Experiment

To create the MLflow experiment in Databricks, navigate to Experiments in the left-hand pane and click Create Experiment under Custom (see Figure 4-4). Name the experiment hotel_bookings_cancellations_prediction and leave the Artifacts location blank, and click Create (Figure 4-5).

Figure 4-4. Creating a traditional experiment

Figure 4-5. Creating an experiment

The experiment’s home page will appear, initially empty since no runs have been performed (Figure 4-6). For large-scale projects, consider using an external database (e.g., MySQL or PostgreSQL) as the MLflow tracking server’s backend to enable centralized tracking across teams and projects.

Figure 4-6. The hotel_bookings_cancellations experiment home page

Model Building

The Model Building notebook, available at /chapter04/04-03-Model Building in the book’s Git repository, guides you through experimenting with machine learning models to predict the binary is_canceled feature, which indicates whether a booking was canceled (1) or not (0). The goal is to model this outcome using various algorithms.

By working through the notebook, you’ll learn to evaluate model performance, predictive accuracy, and other key metrics. For an overview of the notebook’s structure, refer to the table of contents shown in Figure 4-7.

Figure 4-7. The table of contents of the Model Building Notebook

Reading the Delta File

Here, you read the hotel_bookings.delta Delta file and convert it to a pandas dataframe:

# Define the path to your Delta table
delta_file = "dbfs:/FileStore/datasets/hotel_bookings.delta"

# Read the Delta table as a Spark DataFrame
spark_df = spark.read.format("delta").load(delta_file)
# Convert the Spark DataFrame to a Pandas DataFrame
bookings_df = spark_df.toPandas()
# Display the Pandas DataFrame
bookings_df.head()

Prepare our Test and Training Data

With the dataset loaded, the next step is to prepare the data for model training. This involves creating a feature matrix and a target variable, splitting the data into training and testing sets, and scaling the variables to ensure they have a mean of zero and a standard deviation close to one. Finally, you will visualize the prepared dataset to confirm it is ready for effective model training.

Define the Feature Matrix and the Dependent Variable

Next, you separate the feature matrix (X) and the target variable (y) from the DataFrame. In this case our target variable is the is_canceled column:

# Separate the feature matrix (X) and the target variable (y) from the original DataFrame

# Create the feature matrix 'X' by dropping the target column 'is_canceled'
# This matrix contains all the independent variables that will be used to predict the target.
# Each row represents a booking, and each column represents a feature such as 'lead_time', 'country', 'market_segment', etc.
# The 'is_canceled' column is excluded because it represents the dependent variable we want to predict.
X = bookings_df.drop(columns=['is_canceled'])

# Create the target variable 'y' by selecting the 'is_canceled' column from the original DataFrame
# This target variable indicates whether a booking was canceled (1) or not (0).
# 'y' will be used as the dependent variable during model training and evaluation.
y = bookings_df['is_canceled']

Splitting the pandas dataframe into a training and test dataset

To prepare your data for model training and evaluation, you will split the dataset into training and testing subsets using an 80/20 split ratio. The train_test_split function from scikit-learn ensures a random division of the data while maintaining reproducibility with random_state=42.

The training set (80%) is used to train the model, and the testing set (20%) is reserved for evaluating its performance on unseen data.

The test_size=0.2 parameter specifies the split ratio, and random_state=42 ensures consistent results across runs.

After splitting, you can inspect the shapes of the resulting subsets (X_train, X_test, y_train, y_test) to confirm the data distribution.

# Split the data into training and testing sets using an 80/20 split ratio

# The train_test_split function from scikit-learn is used to randomly split the dataset into two subsets:
# - The training set: This subset will be used to train the machine learning model.
# - The testing set: This subset will be used to evaluate the model's performance on unseen data.

# Arguments:
# - 'X' and 'y' are the feature matrix and target variable, respectively.
# - test_size=0.2 specifies that 20% of the data will be set aside for testing, while the remaining 80% will be used for training.
# - random_state=42 ensures that the split is reproducible. By setting a seed value (42), the function will produce the same split every time it is run, which is useful for consistency in model evaluation and comparison.

X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=42)

# Display the record counts
X_train.shape, X_test.shape, y_train.shape, y_test.shape

This produces the following output:

((69588, 62), (17397, 62), (69588,), (17397,))

This shows that we will be using about 70K records for training and 17K records for testing.

Scaling the X variables

Scaling your feature matrix (X) ensures that all independent variables contribute equally to the model, preventing features with larger numeric ranges from disproportionately influencing performance. This process standardizes features to have a mean of 0 and a standard deviation of 1, which is critical for many machine learning algorithms.

Here’s how scaling is performed:

Initialize the Scaler

A StandardScaler object is used to standardize features by removing the mean and scaling to unit variance.

Fit and Transform Training Data

The fit_transform() method calculates the mean and standard deviation of X_train and transforms it accordingly, ensuring a standardized training set.

Transform Test Data

The transform() method applies the same scaling parameters from X_train to X_test, ensuring consistent scaling between training and testing data.

# Step 1: Initialize the StandardScaler object
# StandardScaler standardizes features by removing the mean and scaling to unit variance.
# It transforms the data to have a mean of 0 and a standard deviation of 1, which is
# important for many machine learning models that are sensitive to the scale of input features.
scaler = StandardScaler()

# Step 2: Fit the scaler on the training data and then transform it
# The fit_transform() method first computes the mean and standard deviation of X_train,
# and then transforms X_train by subtracting the mean and dividing by the standard deviation.
# This ensures that X_train is now standardized with mean 0 and variance 1.
X_train_scaled = scaler.fit_transform(X_train)

# Step 3: Use the same scaler to transform the test data
# The transform() method uses the mean and standard deviation computed from the training data
# to standardize X_test. This ensures that the test data is scaled consistently with the training data.
# Note: Do not use fit_transform() on test data, as it would compute different scaling parameters,
# leading to inconsistent scaling and possibly biased model evaluation.
X_test_scaled = scaler.transform(X_test)

Visualizing the Standard Deviation and Mean

To confirm the effectiveness of feature scaling, you will create density plots for the top five features, selected based on variance and correlation with the target variable. This helps visualize the distribution and validate the standardization process.

1. Calculate Variance and Correlation

Compute the variance of each feature in the scaled training dataset to identify features with the highest dispersion.

Calculate the correlation of each feature with the target variable (y_train), identifying those most predictive.

2. Select Top Features

Extract the top five features with the highest variance and the top five with the highest absolute correlation with the target.

Combine these features into a unique list to ensure no duplicates.

3. Generate Density Plots:

Use subplots to visualize the density distribution of the selected features.

Each plot shows the scaled feature distribution, with a vertical line marking the mean.

4. Visualization Layout

Arrange plots dynamically in a grid for clear presentation. Adjust spacing and size to ensure readability.

By examining these density plots, you verify that the scaling process successfully standardized the feature matrix, as indicated by distributions centered around zero with reduced variance.

# Convert X_train_scaled back to a DataFrame with the original feature names
X_train_scaled_df = pd.DataFrame(X_train_scaled, columns=X_train.columns)

# Step 1: Calculate variance and correlation of features
# Calculate variance for each feature in the scaled training dataset
# Variance measures the dispersion of data points. High variance indicates
# that a feature might have a stronger impact on the model.
feature_variance = X_train_scaled_df.var()

# Calculate the correlation of each feature with the target variable
# Correlation measures the linear relationship between a feature and the target.
# Features with high absolute correlation values are more predictive.
feature_correlation = X_train_scaled_df.corrwith(y_train)

# Step 2: Select top 5 features based on variance and correlation
# Select the top 5 features with the highest variance
top_variance_features = feature_variance.nlargest(5).index

# Select the top 5 features with the highest absolute correlation with the target
top_correlation_features = feature_correlation.abs().nlargest(5).index

# Step 3: Combine the selected features into a single list
# Use a set to avoid duplicates, then convert back to a list.
# This ensures that if a feature is selected based on both variance and correlation, it is only included once.
selected_features = list(set(top_variance_features).union(set(top_correlation_features)))

# Step 4: Visualize the selected features using density plots
# Determine the number of features to visualize
num_features = len(selected_features)

# Calculate the number of rows and columns for subplots
# Use 2 columns for visualization, and calculate the number of rows dynamically
rows = (num_features // 2) + 1  # If num_features is odd, add one more row
cols = 2                        # Fixed number of columns

# Set the overall figure size, with height adjusted based on the number of rows
plt.figure(figsize=(16, 4 * rows))

# Loop through each selected feature and create a density plot
for i, feature in enumerate(selected_features):
    # Create a subplot for each feature
    plt.subplot(rows, cols, i + 1)  # Position subplots in a grid layout
    sns.kdeplot(X_train_scaled_df[feature], label='Scaled', color='orange', fill=True)

    # Set the title of each subplot, with adjusted font size and padding to prevent overlap
    plt.title(f'Density Plot of {feature} (Scaled)', fontsize=14, pad=20)

    # Plot a vertical line indicating the mean value of the feature
    plt.axvline(X_train_scaled_df[feature].mean(), color='red', linestyle='--', label='Mean')

    # Set x-axis label for each feature plot
    plt.xlabel(feature, fontsize=12)

    # Display legend on each subplot
    plt.legend()

# Adjust the spacing between subplots to ensure clarity
plt.tight_layout(pad=3.0)  # Add padding to prevent titles and plots from overlapping
plt.show()  # Display the combined plot

Some of the generated density plots are displayed in Figure 4-8.

Figure 4-8. Density plots for the most relevant features.

From the density plots we can conclude that we indeed have a mean of zero, and a low value for the calculated variance, indicating that our scaling process was successful.

Training our first Model

Set the MLflow Experiment

Next, you set the MLflow context to the experiment created earlier. This step ensures that all MLflow runs, including logged metrics, parameters, and artifacts, are associated with the correct experiment for organization and tracking purposes.

The mlflow.set_experiment()function specifies the experiment’s location within the Databricks workspace. In this case, the experiment is named hotel_bookings_cancellations_prediction and is stored in the user’s directory /Users/bhaelen2018@outlook.com/.

By setting this context, MLflow organizes and tracks different model runs and their associated data, facilitating efficient experiment management during the machine learning workflow.

# Set the MLflow experiment
# This command sets the current experiment where all the MLflow runs will be logged.

# The 'experiment_name' argument specifies the location of the experiment within the Databricks workspace.
# In this case, the experiment is named 'hotel_bookings_cancellations_prediction' and is stored
# under the user's directory '/Users/bhaelen2018@outlook.com/' in the Databricks workspace.

# MLflow will use this experiment to organize and track different model runs, metrics, parameters,
# and artifacts (such as models) that you log during the machine learning workflow.
mlflow.set_experiment(experiment_name = '/Users/bhaelen2018@outlook.com/hotel_bookings_cancellations_prediction')

Training a Logistic Regression Model

You are now ready to build and train a logistic regression model while using MLflow to track the experiment’s metrics, parameters, and model artifacts. To streamline this process, you will enable MLflow’s automatic logging feature, which captures all relevant information without requiring explicit log statements. Below is the code for the model:

from mlflow.models import infer_signature

# Step 1 - Enable automatic logging of parameters, metrics, and model for scikit-learn models
mlflow.sklearn.autolog()

# Step 2 - Start a new MLflow run and log the model
with mlflow.start_run(run_name='logistic_regression_model') as run1:

    # Step 3 - Train your model
    model = LogisticRegression(max_iter=10000, class_weight='balanced')
    model.fit(X_train_scaled, y_train)  # Pass y_train directly as a pandas Series

    # Step 4 - Make predictions on the test set
    y_pred_lr = model.predict(X_test_scaled)

    # Step 5 - Calculate performance metrics
    acc_lr = accuracy_score(y_test, y_pred_lr)  # Use y_test directly
    conf_matrix = confusion_matrix(y_test, y_pred_lr)
    clf_report = classification_report(y_test, y_pred_lr)

    # Step 6 - Infer the model signature from the input and output data
    signature = infer_signature(X_train_scaled, y_pred_lr)

    # Step 7 -Log the model with the inferred signature
    model_name = "Logistic_Regression_Model"
    mlflow.sklearn.log_model(model, model_name, signature=signature)

    # Step 8 - Print out the model's performance
    print(f"Accuracy Score of Logistic Regression: {acc_lr:.4f}")
    print(f"Confusion Matrix:\n{conf_matrix}")
    print(f"Classification Report:\n{clf_report}")

    # Step 9 - Access and print the ID of the active MLflow run
    run1 = mlflow.active_run()
    print("Active run ID:", run1.info.run_id)

Key Explanations:

You enable MLflow’s automatic logging with mlflow.sklearn.autolog(). This feature captures critical metrics, parameters, and model artifacts seamlessly, eliminating the need for manual logging.

You start a new MLflow run using the mlflow.start_run() context manager. Assigning the run a descriptive name, such as 'logistic_regression_model', helps you easily organize and identify experiments. The use of a with block ensures that the run closes automatically after executing the enclosed code.

Inside the with block, you initialize a logistic regression model and train it on the scaled training data (X_train_scaled) and target variable (y_train), passed directly as a Pandas Series.

You make predictions on the test dataset (X_test_scaled).

You calculate performance metrics, including accuracy, a confusion matrix, and a classification report. These metrics provide detailed insights into your model’s performance.

Using infer_signature(), you infer the model’s input and output structure, which is crucial for ensuring compatibility during deployment.

You log the model, along with its inferred signature, using mlflow.sklearn.log_model().

You display the model’s critical performance metrics, such as the accuracy score, confusion matrix, and classification report.

You retrieve and print the active MLflow run’s unique ID. This ID is useful for tracking and accessing the logged run in the MLflow UI.

The confusion matrix is a powerful tool you can use to evaluate the performance of your classification model. It shows the counts of correct and incorrect predictions across different classes, helping you understand where your model excels and where it might need improvement. For binary classification problems, the confusion matrix is a 2x2 table that gives you valuable insights into the types of errors made by your model.

Below is an example of a confusion matrix:

              Predicted Class
              +-------------+--------------+
              |   Positive  |   Negative   |
Actual Class  +----------------------------+
Positive      |     TP      |     FN       |
Negative      |     FP      |     TN       |

Here’s how to interpret these terms:

True Positives (TP)

Cases where the model correctly predicted the positive class, such as accurately identifying a booking cancellation.

True Negatives (TN)

Cases where the model correctly predicted the negative class, such as correctly predicting that a booking was not canceled.

False Positives (FP)

Cases where the model incorrectly predicted the positive class (Type I error). For example, predicting a cancellation when there wasn’t one.

False Negatives (FN)

Cases where the model incorrectly predicted the negative class (Type II error). For instance, predicting that a booking was not canceled when it actually was.

Using the confusion matrix, you can better understand your model’s behavior and take targeted steps to improve its performance. If you’ve generated a confusion matrix in the earlier steps, compare it to this explanation to see how your model’s predictions align.

After training your model, you accessed the active MLflow run using mlflow.active_run(). This function retrieved the details of the current run, including its unique ID. This ID is crucial for tracking, querying, or managing the results of a specific run in the MLflow interface.

The full output of your code is shown in Figure 4-9.

Figure 4-9. Output of the logistic regression run.

Once the model training is complete, you can click on the 1 run link in the cell output to view detailed information about this run. This page will show all the metrics, parameters, and model artifacts that were captured during execution.

The displayed parameters for the run are shown in Figure 4-10.

Figure 4-10. Logical Regression Model parameters

For our logistic regression model, MLflow has automatically logged parameters such as fit_intercept=True (whether or not the model should calculate the intercept) and the max_iter=100 (the maximum number of iterations the optimization algorithm will perform during the model’s training process).

The Metrics output of the run are shown in Figure 4-11

Figure 4-11. Model Metrics

Note

The confusion matrix breaks down a model’s predictions into four categories: true positives (TP), true negatives (TN), false positives (FP), and false negatives (FN). Using this breakdown, you can compute key performance metrics such as accuracy and recall.

Accuracy measures the proportion of correct predictions and is calculated as (TP+TN)/(TP+TN+FP+FN). In your case, you have a confusion matrix with 9,281 true positives, 3,673 true negatives, 1,074 false positives, and 3,369 false negatives, the accuracy score would be 74%. Accuracy is a helpful overview metric but may not reflect true performance in imbalanced datasets.

Recall, on the other hand, focuses specifically on the positive class and measures the proportion of actual positives correctly identified by the model. It is calculated as TP/(TP+FN). In the same confusion matrix, the recall would be 74%, showing that the model successfully identified 74% of actual positive cases.

While accuracy provides a general assessment, recall is particularly valuable in scenarios where missing positive cases (false negatives) is costly, such as fraud detection or medical applications, such as cancer detection. Using both metrics together offers a more balanced view of model performance.

If you examine the Metrics section, you will find a number of metrics as shown in Table 4-1: Model Metrics

Table 4-1. Model Metrics

Metric
Description

accuracy_score_X_test_scaled

The accuracy of the model on the test set

training_f1_score

The harmonic mean of precision and recall.

training_precision_score

The proportion of positive predictions that were actually correct.

training_roc_auc

The ability of the model to distinguish between positive and negative classes.

training_accuracy_score

The proportion of correctly classified instances during training

training_log_loss

The performance of a classification model based on the probability it assigns to each class

training_recall_score

Measures the proportion of actual positives that were correctly identified.

training_score

Represents the overall training accuracy or training R² score

These metrics provide a comprehensive view of your model’s performance. For example:

Accuracy indicates the proportion of correct predictions on the test set.

F1-score combines precision and recall into a single metric, which is particularly valuable when the dataset is imbalanced.

Precision focuses on the correctness of positive predictions, while Recall emphasizes the model’s ability to identify all positive cases.

ROC AUC measures the model’s ranking ability, offering additional insights beyond simple accuracy.

Log Loss evaluates the model’s confidence in its predictions, penalizing overconfidence in incorrect predictions.

By examining these metrics, you gain a well-rounded understanding of your model’s performance, enabling better decision-making for model selection and tuning.

In the Tags section, you’ll find that MLflow automatically logs metadata such as estimator_class and estimator_name. These tags provide valuable information about the model type (e.g., a logistic regression model in this case) and its specific configuration during the run. Tags are particularly useful for organizing and filtering runs, especially in large projects where you work with multiple models and need to track their distinct characteristics efficiently.

In the Artifacts section, you’ll discover various output files associated with the trained model. These artifacts include files such as model.pkl (the serialized model), conda.yaml (environment configuration), and requirements.txt (dependency list). Among them, the MLmodel file stands out as it contains critical details about the model, including its input/output schema and available deployment flavors. Figure 4-12 illustrates how these details are displayed within the MLmodel file, offering insights into the model’s structure and readiness for deployment.

Figure 4-12. Artifacts of run 1 of the model.

Flavors are what make MLflow models highly versatile and deployable across different environments. For this model, two flavors are logged:

python_function

This flavor enables the model to be deployed and queried as a Python function, facilitating seamless integration into applications and workflows.

sklearn

This flavor saves the model in a format compatible with Scikit-learn, utilizing a pickle file for serialization. It allows for straightforward loading and prediction using Scikit-learn tools.

Additionally, the requirements.txt file is logged, which lists the Python dependencies and their versions. This file ensures that the model’s environment can be reliably reproduced for consistent performance across different deployments.

With a single line of code (mlflow.sklearn.autolog()), MLflow provides end-to-end tracking for your machine learning model. It captures everything from key metrics and parameters to model artifacts and environment configurations. This automation simplifies the process of managing, comparing, and deploying models, making MLflow an indispensable tool for efficient machine learning workflows.

Additional Model Runs

In the 04-03-Model Building notebook, you will find three additional machine learning models, all of which utilize MLflow’s autologging functionality. While the details of each model run are not covered in this book, you are encouraged to review the notebook to explore the implementation and insights these models provide.

A summary of each model run, including metrics and key details, can be found on the MLflow Experiments page, as shown in Figure 4-13. This summary offers a quick overview and comparison of the runs, allowing you to analyze performance at a glance.

Figure 4-13. Summary of all runs with the experiment.

To compare different runs, simply select the runs you want to analyze and click the Compare button, as shown in Figure 4-14. This feature allows you to easily evaluate the performance and parameters of multiple runs side by side, helping you identify the best-performing model.

Figure 4-14. Comparing multiple Runs

In the Comparing Runs screen, you can evaluate metrics, parameters, start and end times, and the duration of each run, along with a wealth of other details. By scrolling down to the Metrics section, you can click on accuracy to directly compare the accuracies of each model, as demonstrated in Figure 4-15.

Figure 4-15. Comparing the accuracy of the different models

We encourage you to explore the various graphs and tables in the Comparison screen, as they offer a wealth of information and provide an excellent opportunity to deepen your understanding of the different models. Once you’ve thoroughly reviewed the comparisons, you’ll select your preferred model to proceed with predictions.

In this machine learning experiment, you worked with four models to predict cancellations for hotel bookings:

Logistic Regression

K-Neighbors Classifier

Decision Tree Classifier

Random Forest Classifier

After evaluating their performance metrics using MLflow runs, we assume that you selected the logistic regression model as your preferred choice for the remainder of this chapter.

Model Predictions

Let’s now delve into serving predictions using the selected model by loading it directly from an MLflow run. This process demonstrates how to operationalize your model for practical use.

You can find the corresponding code for this section in the /chapter04/04-04-Model Predictions notebook available in the book’s GitHub repository.

Loading the Model from an MLflow run

To begin, retrieve the correct Run ID from the MLflow experiment’s runs screen. Navigate to the hotel_bookings_cancellations_prediction experiment, and select the logistic_regression_model to go to its detail page, as illustrated in Figure 4-16.

Figure 4-16. Selecting the logistic regression model from the runs screen

To load a model from an MLflow run, you need to construct its URI (Uniform Resource Identifier). The URI uniquely identifies the model within the MLflow tracking system. It combines the Run ID of the experiment and the model name assigned during logging. Here’s the code for defining and verifying the model URI:

# The unique identifier for the MLflow run where the model was logged.
# This ID corresponds to the specific experiment run in MLflow.
run_id = "06ab7ac6a53c4431a0db34827220ddf6"

# The name of the model as it was saved during the experiment run.
# This is the name you used when logging the model with MLflow.
model_name = "Logistic_Regression_Model"

# The URI (Uniform Resource Identifier) for the model.
# This URI follows the format: "runs:/<run_id>/<model_name>".
# It is used by MLflow to locate and load the model from the specified run.
model_uri = f"runs:/{run_id}/{model_name}"

# Print the model URI to verify that it is correct.
# This will output the full model URI to the console.
print(f"Model URI: {model_uri}")

The constructed URI points to the specific model artifacts stored in MLflow for a given run. To load the model, you use the mlflow.sklearn.load_model() function, which retrieves the model based on the URI. Once loaded, the model behaves as a standard Scikit-learn object, ready for predictions or evaluations.

# Load the model from the specified model URI in MLflow.
# The model_uri should point to the location where the model was saved during a previous run.
# This function loads the scikit-learn model that was logged using mlflow.sklearn.log_model().
model = mlflow.sklearn.load_model(model_uri=model_uri)

With the model loaded, you can generate predictions on the test data using its predict() method. Ensure that your test features (X_test_scaled) are correctly preprocessed and aligned with the input format expected by the model. The predictions will be returned as an array of class labels (e.g., 0 or 1 for binary classification).

# Use the loaded model to make predictions
y_pred = model.predict(X_test_scaled)

# Count the number of 0s and 1s in the predictions
num_zeros = sum(1 for pred in y_pred if pred == 0)
num_ones = sum(1 for pred in y_pred if pred == 1)

# Print the results
print(f"Number of 0s: {num_zeros:,}")
print(f"Number of 1s: {num_ones:,}")

Since the model is a scikit-learn model, you can leverage the standard scikit-learn API, including the predict() function, to generate predictions. In this example, you print the number of predictions for each class (0 and 1). This step is useful for validating the predictions and later comparing them with outputs generated using other methods of loading the model.

The code produces the following output:

Number of 0s: 10,355
Number of 1s: 7,042

This output shows the class distribution in the test data predictions, helping you verify that the model is working as expected. Comparing these counts across different loading methods ensures consistency and reliability in the model’s behavior.

Using the model as a Spark UDF

MLflow provides the capability to load the model as a Spark User-Defined Function (UDF), making it possible to generate predictions directly on Spark DataFrames. This feature is especially valuable when dealing with large-scale datasets in distributed environments, as it allows you to efficiently perform inference at scale within a Spark pipeline.

To use the model for predictions in a distributed environment, you first load it as a Spark User-Defined Function (UDF) using MLflow’s pyfunc interface. The pyfunc format in MLflow is a versatile representation that supports models from various machine learning frameworks. The code snippet below demonstrates how to create a Spark UDF:

# Create a Spark UDF (User Defined Function) from a model logged in MLflow using the pyfunc interface.
# 'pyfunc' is a generic format in MLflow that supports loading models from various machine learning frameworks.
# This allows the loaded model to be used in a Spark DataFrame as a UDF for making predictions.

# spark: The active Spark session in Databricks or your Spark environment.
# model_uri: The URI of the MLflow model, which points to the location where the model was saved (logged) in MLflow.
# env_manager: Specifies which environment manager to use for the UDF. In this case, 'virtualenv' is used.
# 'virtualenv' tells MLflow to manage the environment using Python's virtual environment system instead of conda.
pyfunc_udf = mlflow.pyfunc.spark_udf(spark, model_uri=model_uri, env_manager='virtualenv')

This function handles the necessary environment setup, ensuring that all dependencies are correctly managed. Next, you will need to create a Spark dataframe from our X_test_scaled pandas dataframe:

# Create a Spark dataframe from our scaled X Test Values
X_test_sp = spark.createDataFrame(X_test_scaled)

You can now use the Spark UDF to generate predictions for the test dataset. The following code demonstrates how to apply the UDF to a Spark DataFrame:

from pyspark.sql.functions import struct

# Apply the UDF (pyfunc_udf) to the Spark DataFrame (X_test_sp) and create a new column 'prediction'.
# 'withColumn()' is a PySpark function that adds a new column or replaces an existing column in a DataFrame.
# 'pyfunc_udf' is the UDF (created earlier using mlflow.pyfunc.spark_udf) which loads the MLflow model and allows us to make predictions.
# 'struct(*(X_test_sp.columns))' combines all columns from the DataFrame (X_test_sp) into a single struct,
# so that the UDF can process them together as model inputs. The '*' operator unpacks the list of column names.

predicted_df = X_test_sp.withColumn('prediction', pyfunc_udf(struct(*(X_test_sp.columns))))
display(predicted_df)

As shown in the code, we add a new column named prediction to the DataFrame, which contains the model’s prediction results. If the prediction output is an array (depending on the model’s output format), we extract the first element using col('prediction')[0]. This ensures the prediction column contains scalar values for easier interpretation. If the model already outputs scalars, this step can be omitted.

To verify the distribution of predictions, the code counts the occurrences of 0 and 1 in the prediction column:

from pyspark.sql.functions import struct, col, expr

# If 'prediction' is an array, extract the first element
# Adjust based on your model output if it's a scalar, this step can be skipped
new_predicted_df = predicted_df.withColumn('prediction', col('prediction')[0])

# Count the number of 0 and 1 predictions
count_zeros = new_predicted_df.filter(col('prediction') == 0).count()
count_ones = new_predicted_df.filter(col('prediction') == 1).count()

# Display the counts
print(f"Number of 0 predictions: {count_zeros:,}")
print(f"Number of 1 predictions: {count_ones:,}")

This prints out:

Number of 0 predictions: 10,355
Number of 1 predictions: 7,042

We observe that the prediction results match exactly with those obtained using the scikit-learn model, confirming consistency and accuracy across both methods.

Loading the model as a generic PyFunc model

As your final option, you can also load the model as a generic PyFunc model, which is particularly useful for deploying models in production due to its flexibility and universal API. The PyFunc interface allows models from various frameworks (e.g., scikit-learn, TensorFlow) to be loaded with a consistent API, making it easier to integrate into diverse environments and applications. This format is ideal when compatibility or seamless integration across platforms is a priority.

The code for loading the model using the PyFunc interface is as follows:

# Import the 'pyfunc' module from MLflow, which provides functionality for loading models
# in a generic Python function (PyFunc) format. This format allows models from various frameworks
# (scikit-learn, TensorFlow, etc.) to be used with a consistent API.
import mlflow.pyfunc

# Load the model from MLflow using the 'load_model' function from the 'pyfunc' module.
# 'model_uri' specifies the location of the model in MLflow (which can be local or remote).
# The model is loaded as a PyFunc model, which provides a universal interface to make predictions
# regardless of the underlying machine learning framework.
model = mlflow.pyfunc.load_model(model_uri=model_uri)

# Make predictions on the scaled test data, which is stored in 'X_test_scaled'.
# 'X_test_scaled' is expected to be a Pandas DataFrame or NumPy array that contains the input features.
# The 'predict()' method is used to generate predictions using the loaded PyFunc model.
y_pred = model.predict(X_test_scaled)

# At this point, 'y_pred' will contain the model's predictions for each row of 'X_test_scaled',
# which could be either class labels (for classification) or continuous values (for regression).

The pyfunc model also provides a predict() method, similar to the scikit-learn model. This makes it easy to generate predictions in a familiar way, regardless of the underlying framework used to build the model. You can double-check the prediction results by counting the number of zeros and ones in the output, as shown in the code below:

# Count the number of 0s and 1s in the predictions
num_zeros = sum(1 for pred in y_pred if pred == 0)
num_ones = sum(1 for pred in y_pred if pred == 1)

# Print the results
print(f"Number of 0s: {num_zeros:,}")
print(f"Number of 1s: {num_ones:,}")

This produces the following output:

Number of 0s: 10,355
Number of 1s: 7,042

These results are identical to those obtained with the previous prediction methods, demonstrating consistency across different model loading and prediction approaches.

Model Deployment

As you continue your journey through machine learning within the Databricks Lakehouse environment, the next step is registering and deploying models using MLflow. This process is essential for transitioning models from development to production, ensuring they are effectively tracked, versioned, and accessible for real-time inference.

In this section, you will focus on moving a model into production. First, you will register the model in the MLflow Model Registry using the MLflow API. Once registered, you will explore the model’s detail page to review its parameters, inputs, outputs, and other key metadata. After thorough testing, you will promote the model to the production stage, making it ready for deployment.

As part of this workflow, you’ll deploy the model to Unity Catalog, Databricks’ recommended solution for governing and deploying machine learning models. Unity Catalog provides centralized access control, auditing, lineage tracking, and model discovery across workspaces, extending these powerful governance features to ML models. Models registered in Unity Catalog are also fully compatible with the open-source MLflow client, ensuring flexibility and ease of integration.

The code for this section is included in the /chapter04/04-05-Model Deployment notebook, available in the book’s GitHub repository.

Creating a Catalog and Schema

Before registering a model in Unity Catalog, it’s important to ensure that both a valid catalog and schema exist for organizing and managing your models. The following SQL commands ensure the required catalog (book_ai_ml_lakehouse) and schema (registered_models) are created:

%sql
-- Make sure the Catalog is created
CREATE CATALOG IF NOT EXISTS `book_ai_ml_lakehouse`;
-- Create the schema for the models
CREATE SCHEMA IF NOT EXISTS book_ai_ml_lakehouse.registered_models;
-- Make sure we can see the newly created schema
SHOW SCHEMAS IN book_ai_ml_lakehouse

The output of the SHOW SCHEMAS command confirms that the registered_models schema has been successfully created alongside other schemas, as shown below:

+-------------------+
|automl             |
|default            |
|information_schema |
|registered_models  |
+-------------------+

Registering the Model in the MLFlow Model Registry

After identifying a model that meets your performance criteria, the next step is to register it in the MLflow Model Registry in preparation for deployment. While this book focuses on programmatic registration using the MLflow API, you can also perform the same task through the Databricks Model Registry UI. Both methods provide robust and seamless ways to track, manage, and deploy models effectively.

To register a model in the MLflow Model Registry, you will use the run_id associated with the latest experiment run that generated the model. Additionally, it’s essential to assign a descriptive model_name to make the model easily identifiable in the registry. Since this model will be registered in Unity Catalog, you will include the catalog name (book_ai_ml_lakehouse) and schema name (registered_models) as part of the full model name. This hierarchical naming structure helps organize and manage registered models effectively within Unity Catalog.

The following code demonstrates how to define these components:

# The unique identifier for the MLflow run where the model was logged.
# This ID corresponds to the specific experiment run in MLflow.
run_id = "06ab7ac6a53c4431a0db34827220ddf6"
# The name assigned to the model. This will be used as a reference for the registered model.
model_name = "hotel_booking_cancel_prediction"
# The catalog name in Unity Catalog, representing the top-level namespace.
catalog_name = "book_ai_ml_lakehouse"
# The schema name within the catalog, typically used for organizing registered models.
schema_name = "registered_models"
# Construct the full Unity Catalog model name, including catalog, schema, and model name.
unity_model_name = f"{catalog_name}.{schema_name}.{model_name}"
# Print the full model location in Unity Catalog.
print(f"Unity Model Name: {unity_model_name}")

This will output the complete Unity Catalog model name:

Unity Model Name: book_ai_ml_lakehouse.registered_models.hotel_booking_cancel_prediction

Next, you need to set the Registry URI to Databricks Unity Catalog to inform MLflow to use Unity Catalog for managing the model registry:

# Set the registry URI to Databricks Unity Catalog.
# This tells MLflow to use Unity Catalog for managing the model registry.
mlflow.set_registry_uri("databricks-uc")

With the run_id and unity_model_name now defined, you are now ready to register your model in the MLflow Model Registry. The mlflow.register_model() function allows you to take the model artifact logged during an experiment run and register it under a specific name in the registry. This step is essential for managing the model lifecycle, including tracking versions and promoting models to different stages, such as staging or production. The code below demonstrates this process:

# Create an instance of the MlflowClient.
# This client is used to interact with the MLflow Model Registry,
# enabling operations such as registering models, managing versions, and transitioning their lifecycle stages.
client = MlflowClient()
# Register the model in MLflow's Model Registry within Unity Catalog.
# The 'register_model()' function associates a model artifact, logged during a specific MLflow run,
# with a registered model name in the Unity Catalog. Unity Catalog provides centralized access control,
# auditing, and lineage for models.

# 'runs:/{run_id}/model' specifies the path to the model artifact:
# - 'run_id' is the unique identifier for the MLflow run where the model was logged.
# - '/model' is the subpath indicating the location of the logged model artifact within the run.

# 'unity_model_name' is the name under which the model will be registered in the Unity Catalog Model Registry.
# This name is unique within the catalog and helps track multiple versions of the model.

# The model will be registered in Unity Catalog, ensuring benefits such as centralized governance,
# lineage tracking, and compatibility with the open-source MLflow client.
model_version = mlflow.register_model(f'runs:/{run_id}/model', unity_model_name)
# 'model_version' holds the metadata of the newly registered model version, including the version number.
# Each time a model with the same name is registered, MLflow automatically increments the version number,
# enabling version control and lifecycle management.

Since this is the first time the model is being registered, it will be assigned version 1.0. This is illustrated in Figure 4-18.

Figure 4-17. Output of registering the model

If a model with the same name already exists in the MLflow Model Registry, MLflow automatically increments the version number and registers the new version. This versioning system ensures that every iteration of the model is systematically tracked, creating a clear record of changes and providing convenient access to earlier versions whenever needed..

Viewing the Model in the Models Page

To view and manage the registered model, navigate to the Models page in the Databricks workspace sidebar and click on Models. Under the Unity Catalog section, locate and select the hotel_booking_cancel_prediction model. The model will appear in the list, as shown in Figure 4-18.

Figure 4-18. Our model displayed in the Registered model page.

When you click on the model name, the detailed view of the model within Unity Catalog will be displayed, as illustrated in Figure 4-19.

Figure 4-19. The registered model in Unity Catalog

Manage Model Lifecycle Stages through Aliases

Model aliases provide a flexible way to assign a named reference to a specific version of a registered model. While no alias is currently assigned, you can use aliases like "Champion" to signify the preferred or actively used model version. For instance, assigning the “Champion” alias to a model version allows you to designate it as the leading or preferred version, even if it is not yet in production. This alias can then be used in workloads that require the model, and when you’re ready to update to a new version, you simply reassign the “Champion” alias to the newer model version, ensuring seamless transitions:

# create "Champion" alias for version 1 of model "prod.ml_team.iris_model"
client.set_registered_model_alias(unity_model_name, "Champion", 1)

If you now return to the models page, you will now see that Version 1 has an alias of Champion, as shown in Figure 4-20.

Figure 4-20. The model with the Champion alias.

With the model successfully registered, we are now ready to move on to leveraging it for inferencing. The next section will focus on using the registered model to generate predictions, enabling practical applications and integrations.

Model Inferencing

After successfully registering the hotel_booking_cancel_prediction model in the MLflow Model Registry and assigning it the Champion alias, the next critical step is deploying the model for real-time serving. This process involves setting up a serving endpoint capable of handling prediction requests through HTTP REST APIs.

Deploying the model in this manner ensures seamless integration with external applications and services, enabling real-time predictions in production environments. In this section, you will also compare the results obtained from the serving endpoint with those generated by a local scikit-learn model. This deployment infrastructure not only allows various systems to interact with the model and retrieve predictions on-demand but also enhances operational efficiency and supports a wide range of practical use cases. The code for the section can be found in the /chapter04/04-06-Model Inferencing notebook in the book GitHub repository.

Enabling Model Serving

To deploy the model, navigate to the Models section in the Databricks workspace sidebar and select the hotel_booking_cancel_prediction model. On the model details page, locate the Serve this model button in the upper-right corner, as illustrated in Figure 4-21.

Figure 4-21. The Use model for inference button in the model details page.

After clicking the Create button, the Create Serving Endpoint dialog box will appear, as illustrated in Figure 4-22. In this dialog, set the endpoint name to cancel_prediction_endpoint. For the compute setting, leave it configured as Small, which is ideal for managing initial traffic while maintaining efficient resource utilization.

Figure 4-22. The Set up model inference dialog.

Click the Create button to begin the deployment process. Keep in mind that this step may take some time, as Databricks needs to provision a new small cluster to support the endpoint. Once the deployment is successfully completed, you will be redirected to the Serving Endpoints page for the cancel_prediction endpoint, as shown in Figure 4-23.

Figure 4-23. The serving endpoint page

Invoking the Model Inferencing with a Python client

To invoke the serving endpoint, you will need a Databricks user token. You can generate this token by clicking on your user icon located in the top-right corner of the Databricks workspace and selecting Settings. Within the Settings page, navigate to the Developer section in the left-hand menu. Here, you will find the option to manage your access tokens, as depicted in Figure 4-24.

Figure 4-24. The Developer Settings section

After clicking Manage under Access Tokens, select Generate new token to create a personal access token. In the Generate new token dialog, provide a meaningful comment, such as Model Inferencing, to help identify the token’s purpose later. Choose a token lifetime that aligns with your security requirements, ensuring it is neither too short nor unnecessarily long. Finally, click Generate to create the new token.

Once generated, securely store the token in a Databricks Secret Scope to avoid exposing it in your code. You can then retrieve and use the token in your notebook by following these steps:

# Import the 'os' module, which provides a way to interact with the operating system.
# The 'os' module allows you to set and retrieve environment variables, among other functionalities.
import os

# Set an environment variable 'DATABRICKS_TOKEN' with the value of your Databricks personal access token.
# This token will be used for authenticating API requests to Databricks services (e.g., for model serving or running jobs).
# The 'os.environ' dictionary allows you to set or access environment variables in the current session.
# Replace with your Key Vault URL
databricks_token = dbutils.secrets.get(scope="book", key="DATABRICKS_TOKEN")
os.environ['DATABRICKS_TOKEN'] = databricks_token

By securely retrieving the token using the dbutils.secrets.get function, you avoid hardcoding sensitive credentials in your codebase. This approach ensures your token remains confidential while enabling seamless authentication for Databricks services, such as model serving or running jobs.

Next, navigate back to the serving endpoint page and locate the Use button in the top-right corner. Click this button to open the Query Endpoint dialog. Within the dialog, switch to the Python tab, where you’ll find the automatically generated Python code for interacting with the endpoint. Copy this code, as illustrated in Figure 4-25, for use in your notebooks or testing environment..

Figure 4-25. The Python code to invoke the serving endpoint.

You can paste the copied code into a cell in your notebook. To test the inference, you will invoke the model directly using scikit-learn, and indirectly through the inference endpoint.

# Define a pipeline with a logistic regression model
# The pipeline consists of a single step: 'classifier', which is the LogisticRegression model.
# The LogisticRegression model is configured with:
# - max_iter=10000: This increases the maximum number of iterations the model will take to converge.
# - class_weight='balanced': This automatically adjusts the class weights to handle class imbalance in the dataset.
pipeline = Pipeline([
    ('classifier', LogisticRegression(max_iter=10000, class_weight='balanced'))
])

# Assume 'X_train_scaled' and 'y_train' are preprocessed training data and labels respectively
# Fit the pipeline (train the model) using the scaled training data
pipeline.fit(X_train_scaled, y_train)

# Specify how many predictions we want to compare (the number of samples from the test set)
num_predictions = 25

# Get predictions from the model served in production using the score_model function
# This function sends a request to the deployed model and returns predictions
# We're sending the first 'num_predictions' samples from 'X_test_scaled' to the served model
served_predictions = score_model(X_test_scaled[:num_predictions])

# Get predictions from the locally trained pipeline model on the same test data
# Here, the pipeline's 'predict()' function is used to get predictions from the logistic regression model
model_evaluations = pipeline.predict(X_test_scaled[:num_predictions])

# Convert the served model predictions to a Pandas Series
# Assuming 'served_predictions' is a dictionary with a key 'predictions' containing the prediction results
# This series is named 'Served Model Prediction' for easier comparison
served_predictions_series = pd.Series(served_predictions['predictions'], name="Served Model Prediction")

# Convert the predictions from the locally trained model (pipeline) into a Pandas Series
# This series is named 'Model Prediction' to distinguish it from the served model's predictions
model_evaluations_series = pd.Series(model_evaluations, name="Model Prediction")

# Create a DataFrame that compares the predictions from the served model and the local model
# The DataFrame has two columns: 'Model Prediction' (local model) and 'Served Model Prediction' (served model)
pred_df = pd.DataFrame({
    'Model Prediction': model_evaluations_series,
    'Served Model Prediction': served_predictions_series
})

# Display the DataFrame with both predictions for comparison
# This will print the DataFrame with the predictions side-by-side for each of the test samples
print(pred_df)

By comparing the results from both the local model prediction and the served model prediction, we can confirm that the outputs are identical.

 # Check if all values match between the two columns
if (pred_df['Model Prediction'] == pred_df['Served Model Prediction']).all():
    print("The columns have the same values.")
else:
    print("The columns have different values.")

Which outputs:

The columns have the same values.

Proving that both columns have indeed the same values. The identical results confirmed that the endpoint is serving inference results accurately.

Summary

In this chapter, you worked through a practical machine learning use case on the Databricks Lakehouse platform, where you leveraged MLflow to manage the end-to-end machine learning lifecycle. The project focused on predicting hotel booking cancellations using the hotel_bookings.csv dataset from Kaggle. Throughout the chapter, you experienced the entire machine learning workflow, including data exploration, feature engineering, model training, deployment, and inference.

You started with Exploratory Data Analysis (EDA) to examine the dataset’s structure, uncover patterns, and validate hypotheses. This step provided critical insights that informed the subsequent feature engineering phase, where you prepared the data for model training by handling missing values, encoding categorical variables, and scaling numerical features. These steps ensured the dataset was optimized for training machine learning models.

Next, you trained a Logistic Regression model using MLflow’s autologging feature, which automatically tracked metrics, parameters, and artifacts. You evaluated the model’s performance using accuracy, precision, recall, F1-score, and confusion matrices, gaining insights into its strengths and weaknesses. All the experimental details were logged to MLflow, enabling you to track and organize your work effectively.

After selecting the best model, you registered it in the MLflow Model Registry, ensuring proper versioning and lifecycle management. You then promoted the model to the Production stage and deployed it as a REST API endpoint for real-time inference. You also explored advanced deployment methods, including serving predictions with Spark UDFs and PyFunc models, which allowed you to handle distributed data and integrate with various environments.

Finally, you validated the model by comparing predictions from the local and served models, confirming that the endpoint provided consistent results. This process demonstrated how seamlessly MLflow supports the transition from development to production in a lakehouse architecture.

Through this chapter, you gained hands-on experience with building, managing, and deploying machine learning models using Databricks and MLflow. You now have a deeper understanding of creating scalable, reproducible workflows, empowering you to implement similar solutions in your own projects.

In the next chapter, we will look into Automated Machine Learning (AutoML) within the Databricks Lakehouse. AutoML simplifies and accelerates the model development process by automating key tasks such as feature engineering, model selection, and hyperparameter tuning. We’ll explore how Databricks AutoML enables both novice and experienced practitioners to rapidly develop high-quality models with minimal manual intervention. This approach not only enhances productivity but also allows for more consistent and scalable machine learning workflows, building upon the foundations we’ve established in this chapter.

Chapter 5. AutoML

A Note for Early Release Readers

With Early Release ebooks, you get books in their earliest form—the author’s raw and unedited content as they write—so you can take advantage of these technologies long before the official release of these titles.

This will be the 5th chapter of the final book. Please note that the GitHub repo will be made active later on.

If you have comments about how we might improve the content and/or examples in this book, or if you notice missing material within this chapter, please reach out to the editor at jbleiel@oreilly.com.

As machine learning continues to grow in importance across various industries, there’s an increasing need to make it accessible to a broader range of professionals, including those without extensive coding experience. Databricks AutoML represents a significant advancement in this direction, democratizing machine learning by enabling users to build and train robust models with little to no code. AutoML simplifies the end-to-end machine learning process, automating tasks from data preparation and preprocessing to model training and evaluation. In this chapter, we will introduce the fundamental concepts of Databricks AutoML and demonstrate how it streamlines workflows for creating regression, classification, and time series forecasting models, empowering users to deliver insights faster and more effectively.

What is AutoML?

In this section, you will discover how Databricks AutoML simplifies the machine learning workflow, making it more accessible and efficient. AutoML (Automated Machine Learning) refers to the process of automating various stages of the machine learning pipeline, including data preprocessing, feature engineering, model selection, hyperparameter tuning, and deployment. By reducing the manual effort required, AutoML enables both novice and experienced data practitioners to build high-performing models without deep expertise in machine learning. With Databricks AutoML, tasks such as handling missing values, selecting and tuning models, and registering models with the Model Registry are automated. All you need to do is specify your dataset and define the prediction target. While AutoML offers options to configure additional settings—such as choosing imputation methods or selecting model frameworks—the heavy lifting is fully managed, allowing you to focus on model insights and business outcomes rather than the technical details of implementation.

In this chapter, you will gain hands-on experience training regression, classification, and forecasting models using both the AutoML user interface and the Python API. You’ll also learn how to deploy and serve models using MLflow’s serving capabilities.

By the end of this sectionchapter, you will be well-equipped to leverage Databricks AutoML for streamlining machine learning workflows, enabling efficient model building, training, and deployment, even without extensive coding expertise.

When you initiate an AutoML run, the system prepares your data for training by performing the necessary preprocessing and transformations to ensure compatibility with various models. AutoML then conducts multiple trials, optimizing and evaluating models across frameworks like scikit-learn, XGBoost (https://xgboost.readthedocs.io/en/stable/), LightGBM (https://lightgbm.readthedocs.io/en/stable/), and others. Each trial is logged as a run within your experiment, and AutoML generates a Python notebook for each, containing the code used in the process. This transparency not only enhances understanding but also allows for customization, enabling you to modify and refine the code as needed.

At the time of writing, Databricks AutoML supports three primary types of models:

Regression Models

Used to predict continuous values, with support for models like decision trees, random forests, linear regression (using stochastic gradient descent), XGBoost, and LightGBM.

Classification Models

Designed to categorize data into discrete classes, AutoML supports decision trees, random forests, logistic regression, XGBoost, and LightGBM.

Forecasting Models

For time series data, AutoML offers support for the Prophet and Auto-ARIMA frameworks.

Databricks AutoML also includes several features that enhance model training and evaluation:

Data Preprocessing

Handles missing values using mean, median, or constant imputation for numeric columns.

Large Dataset Handling

Samples large datasets for initial exploration and efficiently distributes data across cluster nodes for faster training.

Imbalanced Dataset Detection

Identifies and addresses class imbalance in classification problems, applying techniques to rebalance the data before training.

Feature Importance Visualization

Uses Shapley values (https://www.investopedia.com/terms/s/shapley-value.asp#:~:text=In%20game%20theory%2C%20the%20Shapley,other%20to%20obtain%20the%20payoff) to provide model explainability, helping you understand the contribution of each feature to the model’s predictions.

Semantic Type Annotations

Accepts hints about data types (e.g., categorical columns) to apply appropriate transformations during training.

Feature Store Integration

Leverages the Databricks Feature Store for accessing preprocessed features, enhancing consistency and efficiency in model training

Databricks AutoML is continually evolving, with each new release of the Databricks ML runtime introducing additional features and enhancements. It is recommended to consult the official Databricks documentation to stay updated.

Using the AutoML UI to Build a Simple Classification Model

In this section, we will explore one of the most well-known datasets in the field of machine learning: the Kaggle Titanic (Titanic - Machine Learning from Disaster | Kaggle) dataset.

The Titanic dataset contains detailed information about the passengers aboard the RMS Titanic, which tragically sank during its maiden voyage in 1912. It provides various attributes, such as age, gender, passenger class, number of siblings/spouses aboard, and fare paid, among others. The objective of this use case is to predict a key outcome, in this case whether a passenger survived or not, based on these features.

Due to its historical context, structured format, and balanced mix of numerical and categorical data, the Titanic dataset serves as an ideal learning tool. It not only highlights the nuances of preprocessing and feature engineering but also offers valuable insights into how AutoML can handle complex tasks like missing data, categorical encoding, and hyperparameter tuning to build robust predictive models

Downloading the Dataset from Kaggle

Please refer to the notebook /chapter05/05-01-Download Titanic Dataset in the book’s GitHub repository. This notebook covers the following key steps:

Ensure the installation of required libraries

Verifies that the necessary libraries, specifically kaggle and kagglehub, are installed and configured correctly.

Download the Titanic dataset

Retrieves the dataset in CSV format directly from Kaggle.

Data preparation and conversion

Reads the CSV file, converts each column to the appropriate data type, and saves the processed DataFrame as a Delta Lake table for efficient storage and subsequent analysis.

Using the AutoML User Interface

With the Titanic dataset now successfully loaded into the Delta table, we are ready to launch our first AutoML experiment. Navigate to the Experiments main page in your environment, where you’ll find three primary options for AutoML: Forecasting, Classification, and Regression. Since our objective is to predict whether a passenger survived—a classification problem—we’ll choose the Classification option to initiate the model training process, as depicted in Figure 5-1.

Figure 5-1. Create AutoML Classification Experiment on the Experiments home page

To begin, click on Start Training to navigate to the Configure Classification Experiment page. At the top of this page, you’ll find a clear, four-step workflow designed to guide you through the experiment creation process seamlessly:

1. Configure the Experiment

Define the experiment parameters by selecting the target column (Survived) and choosing an evaluation metric (e.g., accuracy, precision, or recall) to measure the model’s performance.

2. Join the Features

Integrate and preprocess features from the Delta table. This step ensures that all relevant columns are appropriately prepared and ready for model training, including handling missing values and data types.

3. Train the Model

Leverage AutoML to explore various machine learning algorithms, perform automated feature selection, and optimize hyperparameters to build a robust classification model.

4. Evaluate the Model

Analyze the performance of the trained model using metrics like accuracy, precision, recall, and F1-score. Identify areas for improvement and determine the model’s readiness for deployment.

By following these steps, you can efficiently build a predictive model for the Titanic dataset with minimal manual intervention.

On the right-hand side of the Configure Classification Experiment page, you’ll find a detailed list of the available Databricks runtime versions, along with the specific AutoML features supported by each version. This interface provides valuable insights into the capabilities and enhancements of different runtimes, helping you choose the most suitable environment for your AutoML workflows. As shown in Figure 5-2, each runtime version highlights key AutoML features, such as support for additional algorithms, advanced preprocessing methods, or enhanced model explainability.

Reviewing this listing ensures that you select the runtime version best suited to your AutoML experiments, allowing you to leverage the latest features and optimizations. This step is particularly crucial as Databricks frequently updates its ML runtimes to include new AutoML functionalities and performance improvements. Staying informed about these updates not only enhances your model’s performance but also ensures compatibility with the latest machine learning frameworks and tools.

Figure 5-2. AutoML features in each runtime.

The next step in configuring your AutoML experiment involves providing three key pieces of information. The first requirement is to specify the cluster where the training will be executed. Select the cluster you set up in Chapter 2, ensuring it is configured with the appropriate ML runtime. Keep in mind that only clusters equipped with the ML runtime will appear in the dropdown menu, as this environment is essential for running AutoML processes. Selecting the correct cluster ensures seamless execution and access to the tools necessary for training your model effectively.

The second step is to specify the input training dataset. For this example, select the Titanic table that was prepared earlier and saved as a Delta table. This dataset serves as the foundation for the AutoML experiment, providing both the features and the target column needed for model training. Figure 5-3 illustrates the interface for selecting the input table, making it easy to identify and confirm the dataset to be used for the experiment.

Figure 5-3. Selecting the training table.

Click Select to choose the Titanic table as your input dataset.

The next step in configuring your AutoML experiment is to define the prediction target. For this example, the target variable is the survived attribute, which indicates whether a passenger survived the Titanic disaster. You will also need to provide a name for your experiment. Set the experiment name to titanic_survivors to clearly reflect the purpose of this analysis.

After selecting the prediction target, review the dataset schema displayed on the right-hand side of the interface, as shown in Figure 5-4. This schema view offers a detailed overview of the column names and their data types. Examining the schema at this stage allows you to verify that all necessary columns are included and properly formatted, ensuring the dataset is ready for AutoML training.

Figure 5-4. The schema of the selected dataset.

On this screen, you can confirm that the Survived column has been correctly selected as the prediction target, indicated by the checkmark next to it. At this stage, you can choose to include or exclude specific columns from your training dataset, giving you full control over which features will be utilized for model development. Additionally, you can configure the imputer to handle any missing values in your dataset. For numeric columns, you have four imputation options to choose from:

Mean

Replaces missing values with the column’s average value.

Median

Uses the column’s median value to fill in missing entries, which is robust to outliers.

Most Frequent

Fills missing values with the mode, the most commonly occurring value in the column.

Constant

Allows you to specify a fixed value to replace all missing values.

For string columns:

Most Frequent

Fills missing values with the most frequently occurring value in the column.

Constant

Replaces missing values with a custom value that you define.

We will leave the Imputer setting on Auto, allowing the system to automatically select the most appropriate imputation strategy based on the data type and characteristics of the dataset.

In the Advanced Configuration section, you can customize additional settings tailored to your specific ML problem type. The screen for our classification setup is illustrated in Figure 5-5, listing the available options for fine-tuning your experiment.

Figure 5-5. Advanced configuration for the Configure AutoML Experiment screen

You’ll see in the figure that we choose the F1 score as the primary evaluation metric. The F1 score provides a balanced measure of precision and recall, making it especially suitable for imbalanced datasets like the Titanic survival dataset. During model development, AutoML will leverage three powerful frameworks (LightGBM, Scikit-learn, and XGBoost) to explore a variety of algorithms and identify the most effective model.

By default, the AutoML experiment is configured to run for 120 minutes, but given the simplicity of the Titanic dataset, we will reduce the duration to 20 minutes. This shorter runtime is sufficient for the system to iterate through different algorithms and perform hyperparameter tuning across the specified frameworks. AutoML will automatically select and fine-tune models, ensuring the best-performing model is identified within the allocated time.

Since the Titanic dataset does not include a time column, the field for specifying the training, validation, and testing split will be left blank. For the classification target, we will set the positive label to “1”, which represents passengers who survived the Titanic disaster.

For storing intermediate results, you have the option to use MLflow Artifact storage or DBFS (Databricks File System). We will retain the default option, which is MLflow Artifact storage. Additionally, since a custom data split is not required, this field will also be left blank.

Running the AutoML experiment

With our AutoML experiment fully configured, we are now ready to train a classification model using the F1 Score as the evaluation metric. To begin, click the Start AutoML button. The experiment is set to run for 20 minutes. Once training begins, you’ll see a countdown timer indicating the remaining time, starting at 20 minutes.

Before proceeding with the model training phase, Databricks AutoML performs a critical step: data exploration. This preliminary phase is essential for understanding the dataset’s characteristics, detecting potential issues, and informing key modeling decisions. AutoML automatically generates a data exploration notebook, offering detailed insights into the dataset, which can guide both the automated workflow and any manual adjustments you may choose to implement. As shown in Figure 5-6, a link to this generated notebook is provided for easy access.

Figure 5-6. The automatically generated data exploration notebook.

It is highly recommended that you open the generated notebook to review the various actions taken by AutoML. You’ll notice that these steps closely resemble the exploratory data analysis (EDA) process we performed in Chapter 3.

As the training progresses, you can refresh the experiment page to observe new model runs being logged in real time. Once the training and evaluation phases are complete, the experiment page will display all completed runs. To identify the best-performing models, I added the test_f1_score column to the display grid and sorted the results in descending order by this metric, as illustrated in Figure 5-7.

Figure 5-7. Completed runs sorted by f1 score.

As anticipated, a high number of the runs achieved a perfect F1 score, highlighting the simplicity of our classification problem. You’ll notice a View Notebook for the Best Model button on top of the runs listing. Clicking this button opens a notebook that provides a detailed overview of the training process for the top-performing model. In this instance, the best model was identified as a Decision Tree Classifier, as shown in Figure 5-8.

Figure 5-8. The best performing model notebook leverages a Decision Classifier.

You are encouraged to explore this notebook at your convenience, as it provides valuable insights into the model training process. Additionally, you can customize and modify the notebook to align with your specific requirements. In the next section we will walk through the most critical steps performed by this AutoML-generated notebook.

Loading the Data

The notebook begins by loading the dataset from MLflow storage, where all artifacts are stored. This is consistent with the configuration specified during the experiment setup, ensuring that the dataset is retrieved directly from the MLflow artifact storage for seamless access and reproducibility.

Select the Supported Columns

AutoML automatically identifies and selects columns that are compatible with the machine learning algorithms. It filters out unsupported data types, ensuring only features such as booleans, numerical values, and categorical variables are included in the modeling process. This step ensures compatibility and optimizes the dataset for training.

Preprocessing the Data

AutoML organizes features into distinct groups and applies tailored preprocessing steps to prepare the data for modeling:

Boolean Columns

Boolean values are converted to 0 or 1 for model compatibility, with any missing values imputed to maintain data integrity.

Numerical Columns

Missing values in numerical features are imputed using the mean, and the data is scaled to ensure all features are on a consistent scale, enhancing model performance.

Categorical Columns

Low-cardinality categorical features (those with fewer unique values) are encoded using one-hot encoding, preserving both simplicity and interpretability.

Train Validation Split

AutoML divides the dataset into three distinct sets to optimize model development and evaluation:

Training Set

Used to train the model by fitting it to the data.

Validation Set

Used during hyperparameter tuning to evaluate the model’s performance and guide adjustments.

Test Set

Reserved for assessing the final model’s performance, simulating real-world deployment scenarios and ensuring unbiased evaluation.

Train the Classification Model

The core of this AutoML pipeline involves training a Decision Tree Classifier, which serves as the primary model for the classification task.

Define the Objective Function

AutoML establishes the objective function to evaluate the model’s performance. For classification tasks, this function is based on metrics such as accuracy, F1-score, or AUC-ROC, ensuring the model’s effectiveness is measured appropriately.

Configure the HyperParameter Search Space

AutoML automates the hyperparameter tuning process by defining a search space for key hyperparameters. The notebook outlines the specific search space used, enabling systematic exploration to optimize model performance:

space = {
  "criterion": "entropy",
  "max_depth": 12,
  "max_features": 0.7082296053471858,
  "min_samples_leaf": 0.27057539844789635,
  "min_samples_split": 0.17907363946089483,
  "random_state": 642664015,
}

This search space represents a carefully defined configuration that balances the flexibility and generalization of the decision tree. Parameters like max_depth, min_samples_leaf, and min_samples_split help control the complexity of the model, preventing overfitting. Meanwhile, criterion, max_features, and random_state ensure that splits are both efficient and reproducible. Together, these settings allow AutoML to systematically search for the best combination of hyperparameters to enhance the model’s predictive performance.

Run Trials

AutoML executes multiple trials, systematically testing different combinations of hyperparameters to optimize the model’s performance. Through this iterative process, it continuously refines the model until the best configuration is identified. The results of these trials are displayed in Figure 5-9.

Figure 5-9. The result of the trail runs.

As you can see, we have a perfect f1 score with zero false positives and false negatives.

The created pipeline graph shown in figure 4-9 shows a nice illustration of the pipeline with the different steps:

On the left, we see that for boolean values, a OneHotEncoder was used to encode a “0” or “1” value.

For numerical values a SimpleImputer was used to impute a default value

For categorical values a OneHotEncoder was used to encode the values

Finally, all features were fed into the DecisionTreeClassifier instance.

Patch Pandas version in the logged model.

During the model registration step, AutoML ensures compatibility by patching the pandas version in the logged model artifacts. This is crucial for maintaining consistency during inference, as different environments might have varying dependencies.

Feature Importance

After training, AutoML generates a feature importance report, highlighting which features had the greatest impact on the model’s predictions. This aids in interpretability and helps data scientists understand which variables drive the outcomes.

Inference

The trained model is ready for inference, where it can make predictions on new data. This step represents the transition from model development to practical application. Note that this section is shown as markup text in the notebook, but you can easily move it to a notebook cell. The code would register the model in the model registry, and move the model through the various stages of your MLOps workflow.

Model evaluation Confusion Matrix, ROC, and Precision-Recall Curves

AutoML evaluates the model using key metrics to ensure reliability:

Confusion Matrix:

Provides a summary of true positives, true negatives, false positives, and false negatives, helping to understand misclassifications.

ROC Curve

Plots the true positive rate against the false positive rate, offering a visual representation of the trade-off between sensitivity and specificity.

Precision-Recall Curve

Focuses on the balance between precision and recall, especially useful for imbalanced datasets.

These evaluation metrics are initially applied to the validation data and then to the test set to confirm model performance consistency.

As a next step, you can tweak the notebook to your needs, register the model and serve predictions from the model as you did in chapter 4.

Using the AutoML python API to build a regression model

In the previous section, you explored how to leverage the Databricks AutoML user interface (UI) to quickly and efficiently train machine learning models. While the UI serves as an excellent starting point with its ease of use and intuitive design, there are scenarios where you may need more control and flexibility. For these situations, the ability to programmatically script your experiments becomes invaluable. In this section, you’ll transition to using the AutoML Python API, which allows you to configure and train a regression model directly within a Python notebook, giving you a more hands-on approach to model development.

The Databricks AutoML Python API replicates the capabilities of the UI while providing you with greater flexibility and control over your experiments. By using code, you can initiate and configure AutoML runs directly within your notebook, opening up new possibilities to:

Automate workflows

Seamlessly integrate AutoML into your data pipelines or CI/CD processes, enabling end-to-end automation that accelerates the entire machine learning lifecycle..

Customize configurations

Programmatically adjust parameters and hyperparameters, allowing you to fine-tune the model training process according to your specific requirements and data characteristics.

Enhance reproducibility

Use version control to manage your code, making it easier to share notebooks with your team. This ensures that experiments can be consistently reproduced and reviewed, promoting collaboration and transparency..

The concepts and configuration parameters you’ve explored so far with the AutoML UI are equally applicable when working with the AutoML Python API. Whether you’re training regression, classification, or forecasting models, the API maintains consistent functions and arguments, ensuring a smooth and familiar workflow. This consistency not only streamlines your model-building process but also allows you to transition seamlessly between the UI and programmatic approaches as your needs evolve.

Preparatory Steps

In this section, you will be working with the California Housing Prices (https://www.kaggle.com/datasets/kallolnath1/california-housing-prices-dataset) dataset from Kaggle. Before diving into the AutoML Python API, there are a few preparatory steps you need to complete:

First, you’ll need to download the dataset. To do this, use the notebook labeled 05-02-Download California Housing Prices Dataset. This notebook is designed to handle the entire process, from downloading the dataset to adjusting data types, and ultimately writing the resulting DataFrame to the book_ai_ml_lakehouse.automl.ca_housing_prices Unity table.Feel free to explore this notebook on your own.

Next, you’ll perform exploratory data analysis (EDA) on the dataset using the 05-03-EDA and Cleansing of the CA Housing Dataset notebook. This notebook guides you through analyzing, visualizing, and cleansing the data to prepare it for model training. Once completed, it will write the cleansed features to the: book_ai_ml_lakehouse.automl.ca_housing_prices_features Unity table. As with the previous notebook, you are encouraged to explore this one at your own pace to better understand the EDA and data preparation process.

Using the AutoML API

The California Housing Dataset is used for predicting housing prices in California. Specifically, the dataset provides various socio-economic and geographical factors that influence housing prices, and the goal is to predict the median house value for California districts.

To follow along, open the notebook titled 05-04-CA Housing Dataset Regression with AutoML.

The notebook begins by reading the ca_housing_prices_features Delta table created during the exploratory data analysis (EDA) phase. When you display a sample of the data, you’ll notice that the dataset has already been cleansed. Additionally, the ocean_proximity column has been one-hot encoded, as illustrated in Figure 5-10. This ensures that the data is well-prepared for model training, enabling a seamless transition to the next steps of the AutoML workflow.

Figure 5-10. The California housing dataset, with the ocean_proximity column one-hot encoded.

Once you have loaded the dataset, you perform a train/test split of the dataframe:

# Split the pandas DataFrame into training and testing sets
# 'train_test_split()' is used to split the data into random train and test subsets
# 'df_housing_prices' is the original DataFrame to be split
# 'test_size=0.2' specifies that 20% of the data should go into the test set, while 80% goes into the training set
# 'random_state=42' ensures that the split is reproducible, so you get the same train-test split each time you run it
train_df, test_df = train_test_split(df_housing_prices, test_size=0.2, random_state=42)

Next, you’ll need to convert the Pandas training DataFrame (train_df) to a Spark dataframe. This conversion is necessary because the automl.regress() method expects a Spark dataframe as its first parameter:

# Convert the Pandas Dataframe to a Spark Dataframe,
# since the AutoMLRegressor class expects a Spark Dataframe
train_df_spark = spark.createDataFrame(train_df)

With the data preparation complete, you’re now ready to initiate the AutoML regression:

# Run AutoML regression on the training Spark DataFrame
# 'automl.regress()' is an AutoML function that automatically finds the best regression model
# for the given dataset and target column
# 'train_df_spark' is the input Spark DataFrame containing the training data
# 'target_col='median_house_value'' specifies the target column to predict (dependent variable)
# 'timeout_minutes=10' sets a timeout of 10 minutes, after which the AutoML process will stop,
# even if it hasn't tried all possible models
summary = automl.regress(
    train_df_spark,
    target_col='median_house_value',
    timeout_minutes=10
)

# 'summary' will contain information about the best model found by AutoML, along with metrics,
# feature importance, and other details that can be used for evaluation and interpretation

In this step, you set the target column to median_house_value. Given that this is a relatively small dataset, you set the timeout to just 10 minutes to allow for quicker model exploration.

It’s worth noting that you didn’t explicitly specify an experiment when calling the regress() method. However, AutoML automatically creates an experiment for you. To view it, navigate to the Experiments menu in the left-hand navigation panel. By selecting the most recent experiment from the list, you can access the experiment summary and review all of the individual runs. Once the experiment completes, you add the test R² score to the displayed columns and sort the results, allowing you to quickly identify the best-performing model at the top of the list, as illustrated in Figure 5-11.

Figure 5-11. The median house values experiment with the best model runs

The best model achieved a 0.79 training R² score, indicating a strong fit on the training data. You can confirm this result by reviewing the output of the regress() method in the notebook, as shown in Figure 5-12.

Figure 5-12. The metrics for the best trail.

On top of the results table you see the link to the best trail notebook. After clicking this link, you can explore all the different steps taken by AutoML to create the model.

Analysis of the Best Trial Notebook

This notebook represents the best-performing AutoML-generated model selected for predicting median house prices in California. Databricks AutoML automates the entire machine learning pipeline, from data preprocessing and feature selection to model training, hyperparameter tuning, and deployment. The chosen model, a LightGBM Regressor, demonstrated the highest performance based on validation metrics.

This review provides an in-depth walkthrough of the notebook, highlighting how AutoML optimized the workflow to ensure the most effective model selection. We will explore the key steps in the notebook, including data loading, feature selection, preprocessing, model training, hyperparameter tuning, and final model registration in MLflow for deployment.

By the end of this review, you will gain a clear understanding of how Databricks AutoML streamlines machine learning, ensuring efficient and reproducible model development while minimizing the need for manual intervention.

Loading the Dataset

In this section, Databricks AutoML loads the dataset from MLflow artifacts using the download_artifacts API. The dataset is stored in Parquet format, a commonly used storage format for efficient data retrieval and processing. Once downloaded, the data is read into a Pandas DataFrame, allowing for further transformations and model training. The following code snippet demonstrates how the dataset is retrieved and loaded:

# Download the artifact and read it into a pandas DataFrame
input_data_path = mlflow.artifacts.download_artifacts(run_id="901549f74fd74530a2dc75d530145533", artifact_path="data", dst_path=input_temp_dir)

df_loaded = pd.read_parquet(os.path.join(input_data_path, "training_data"))

Feature Selection

Next, the notebook performs feature selection by selecting only the columns that are supported for model training. AutoML automatically filters out unsupported features and ensures that only relevant attributes are included in the modeling process.

Additionally, categorical features, such as ocean_proximity, have already been one-hot encoded, meaning they are transformed into numerical representations suitable for machine learning models. The following code snippet demonstrates how AutoML selects the final set of features for training:

from databricks.automl_runtime.sklearn.column_selector import ColumnSelector
supported_cols = ["latitude", "ocean_proximity_NEAR_BAY", "median_income", "housing_median_age", "households", "ocean_proximity_NEAR_OCEAN", "ocean_proximity_ISLAND", "ocean_proximity_less_than_1H_OCEAN", "ocean_proximity_INLAND", "longitude", "population", "total_bedrooms", "total_rooms"]
col_selector = ColumnSelector(supported_cols)

Data Preprocessing

In this section, the notebook defines a preprocessing pipeline to handle boolean (categorical) features and numerical columns.

Boolean Columns

In this section, the notebook defines a preprocessing pipeline to handle boolean (categorical) features in the dataset using scikit-learn’s Pipeline and ColumnTransformer objects. Preprocessing is a crucial step in machine learning, ensuring that data is properly formatted and prepared before model training.

The pipeline is designed to:

Convert categorical boolean features into numerical form for compatibility with machine learning models.

Apply transformations systematically, ensuring consistency across training and inference.

Handle missing values and type conversions where necessary.

To set up this pipeline, the notebook first imports the necessary modules:

First, all relevant modules are imported:

from sklearn.compose import ColumnTransformer
from sklearn.impute import SimpleImputer
from sklearn.pipeline import Pipeline
from sklearn.preprocessing import FunctionTransformer
from sklearn.preprocessing import OneHotEncoder as SklearnOneHotEncoder

Next, the notebook sets up a Boolean Pipeline, which is responsible for processing categorical boolean features in the dataset. Since machine learning models require numerical inputs, boolean features must be transformed appropriately before training.

This pipeline performs the following key steps:

Casts boolean features to object type to ensure compatibility with encoding techniques.

Handles missing values (though no imputation is applied in this step).

Applies one-hot encoding to convert categorical boolean features into a format suitable for machine learning.

The following code snippet defines the Boolean Pipeline:

bool_imputers = []

bool_pipeline = Pipeline(steps=[
    ("cast_type", FunctionTransformer(lambda df: df.astype(object))),
    ("imputers", ColumnTransformer(bool_imputers, remainder="passthrough")),
    ("onehot", SklearnOneHotEncoder(handle_unknown="ignore", drop="first")),
])

Finally, the notebook defines the boolean transformers, which apply the Boolean Pipeline to specific categorical features in the dataset. These transformations ensure that boolean-like categorical variables are properly encoded for use in the machine learning model.

In this case, the ocean_proximity feature, which originally contained categorical text values (e.g., "NEAR BAY“, "INLAND“), has already been one-hot encoded into multiple boolean columns. The Boolean Transformer is responsible for:

Applying the previously defined Boolean Pipeline to these columns.

Ensuring categorical variables are encoded consistently before model training.

The following code snippet defines the Boolean Transformer:

bool_transformers = [("boolean", bool_pipeline, [
          "ocean_proximity_NEAR_BAY",
          "ocean_proximity_NEAR_OCEAN",
          "ocean_proximity_ISLAND",
          "ocean_proximity_less_than_1H_OCEAN",
          "ocean_proximity_INLAND"])]

Numerical Columns

This code defines a numerical feature preprocessing pipeline using scikit-learn’s Pipeline and ColumnTransformer. The pipeline is designed to prepare numerical features for machine learning by performing three key transformations: type conversion, missing value imputation, and feature scaling. First, a SimpleImputer fills in missing values with the mean of each respective column, ensuring that no missing data disrupts the model training process. Next, a FunctionTransformer ensures that all numerical columns are correctly formatted as numeric data, preventing potential type-related issues. Finally, a StandardScaler is applied to normalize the numerical features by ensuring they have a mean of 0 and a standard deviation of 1, which helps machine learning models converge faster and perform better.

The pipeline is applied to various demographic, economic, and geographic features of the dataset, such as latitude, longitude, median income, population, and total rooms. Additionally, one-hot encoded categorical variables related to ocean proximity are treated as numerical inputs and processed in the same manner. By automating these transformations, the pipeline ensures data consistency and improves model performance, making it well-suited for use in an AutoML workflow. This approach not only streamlines preprocessing but also enhances reproducibility, ensuring that the same transformations are applied consistently across training and inference phases.

from sklearn.compose import ColumnTransformer]
from sklearn.impute import SimpleImputer
rom sklearn.pipeline import Pipeline
from sklearn.preprocessing import FunctionTransformer, StandardScaler
num_imputers = []
num_imputers.append(("impute_mean", SimpleImputer(), ["households", "housing_median_age", "latitude", "longitude", "median_income", "ocean_proximity_INLAND", "ocean_proximity_ISLAND", "ocean_proximity_NEAR_BAY", "ocean_proximity_NEAR_OCEAN", "ocean_proximity_less_than_1H_OCEAN", "population", "total_bedrooms", "total_rooms"]))
numerical_pipeline = Pipeline(steps=[
    ("converter", FunctionTransformer(lambda df: df.apply(pd.to_numeric, errors='coerce'))),
    ("imputers", ColumnTransformer(num_imputers)),
    ("standardizer", StandardScaler()),
])
numerical_transformers = [("numerical", numerical_pipeline, ["latitude", "ocean_proximity_NEAR_BAY", "median_income", "housing_median_age", "households", "ocean_proximity_NEAR_OCEAN", "ocean_proximity_ISLAND", "longitude", "ocean_proximity_INLAND", "ocean_proximity_less_than_1H_OCEAN", "population", "total_bedrooms", "total_rooms"])]

from sklearn.compose import ColumnTransformer
transformers = bool_transformers + numerical_transformers
preprocessor = ColumnTransformer(transformers, remainder="passthrough", sparse_threshold=0)

Data Splitting

This code handles the train-validation-test split of the dataset, which was automatically performed by Databricks AutoML. During the AutoML process, a special column named _automl_split_col_0000 was added to indicate which records belong to the training, validation, and test sets. The dataset is then filtered based on this column to create separate training (split_train_df), validation (split_val_df), and test (split_test_df) subsets. This approach ensures that data leakage is prevented, and model evaluation is performed on unseen data, maintaining robustness and generalization.

After splitting the dataset, the target column (target_col) is separated from the features for each set. The features (X_train, X_val, X_test) are extracted by removing both the target variable and the split indicator column (_automl_split_col_0000), while the target labels (y_train, y_val, y_test) are stored separately. This structure ensures that the dataset is ready for model training, hyperparameter tuning, and evaluation, following best practices for supervised learning:

# AutoML completed train - validation - test split internally and used _automl_split_col_0000 to specify the set
split_train_df = df_loaded.loc[df_loaded._automl_split_col_0000 == "train"]
split_val_df = df_loaded.loc[df_loaded._automl_split_col_0000 == "validate"]
split_test_df = df_loaded.loc[df_loaded._automl_split_col_0000 == "test"]

# Separate target column from features and drop _automl_split_col_0000
X_train = split_train_df.drop([target_col, "_automl_split_col_0000"], axis=1)
y_train = split_train_df[target_col]

X_val = split_val_df.drop([target_col, "_automl_split_col_0000"], axis=1)
y_val = split_val_df[target_col]

X_test = split_test_df.drop([target_col, "_automl_split_col_0000"], axis=1)
y_test = split_test_df[target_col]

Training the Regression Model

This code defines an AutoML-driven model training and evaluation pipeline using MLflow, LightGBM, and Hyperopt. The objective is to optimize a LightGBM Regressor (LGBMRegressor) through hyperparameter tuning while leveraging early stopping and automatic logging for experiment tracking. The process starts by setting up a separate preprocessing pipeline (pipeline_val) for the validation dataset, ensuring that transformations applied during training are consistently used in validation. MLflow’s autologging is temporarily disabled, and the dataset is preprocessed before model training.

The objective function encapsulates the training, evaluation, and logging workflow for hyperparameter optimization. Within each MLflow experiment run, a new model instance is created using the provided hyperparameters. The LightGBM regressor is trained, incorporating callbacks for early stopping and evaluation logging. After training, MLflow’s evaluation API is used to compute metrics for training, validation, and test sets, which are logged for performance tracking. The function then returns a loss value (negative R² score on validation data), enabling Hyperopt’s optimization algorithm (fmin) to iteratively find the best-performing model. This structured approach ensures automated experimentation, performance tracking, and efficient hyperparameter tuning.

import mlflow
from mlflow.models import Model, infer_signature, ModelSignature
from mlflow.pyfunc import PyFuncModel
from mlflow import pyfunc
import sklearn
from sklearn import set_config
from sklearn.pipeline import Pipeline
from hyperopt import hp, tpe, fmin, STATUS_OK, Trials

# Create a separate pipeline to transform the validation dataset. This is used for early stopping.
pipeline_val = Pipeline([
    ("column_selector", col_selector),
    ("preprocessor", preprocessor),
])

mlflow.sklearn.autolog(disable=True)
pipeline_val.fit(X_train, y_train)
X_val_processed = pipeline_val.transform(X_val)

def objective(params):
  with mlflow.start_run(experiment_id="4009150997045188") as mlflow_run:
    lgbmr_regressor = LGBMRegressor(**params)

    model = Pipeline([
        ("column_selector", col_selector),
        ("preprocessor", preprocessor),
        ("regressor", lgbmr_regressor),
    ])

    # Enable automatic logging of input samples, metrics, parameters, and models
    mlflow.sklearn.autolog(
        log_input_examples=True,
        silent=True,
    )

    model.fit(X_train, y_train, regressor__callbacks=[lightgbm.early_stopping(5), lightgbm.log_evaluation(0)], regressor__eval_set=[(X_val_processed,y_val)])

    # Log metrics for the training set
    mlflow_model = Model()
    pyfunc.add_to_model(mlflow_model, loader_module="mlflow.sklearn")
    pyfunc_model = PyFuncModel(model_meta=mlflow_model, model_impl=model)
    training_eval_result = mlflow.evaluate(
        model=pyfunc_model,
        data=X_train.assign(**{str(target_col):y_train}),
        targets=target_col,
        model_type="regressor",
        evaluator_config = {"log_model_explainability": False,
                            "metric_prefix": "training_"  }
    )
    # Log metrics for the validation set
    val_eval_result = mlflow.evaluate(
        model=pyfunc_model,
        data=X_val.assign(**{str(target_col):y_val}),
        targets=target_col,
        model_type="regressor",
        evaluator_config= {"log_model_explainability": False,
                           "metric_prefix": "val_"  }
   )
    lgbmr_val_metrics = val_eval_result.metrics
    # Log metrics for the test set
    test_eval_result = mlflow.evaluate(
        model=pyfunc_model,
        data=X_test.assign(**{str(target_col):y_test}),
        targets=target_col,
        model_type="regressor",
        evaluator_config= {"log_model_explainability": False,
                           "metric_prefix": "test_"  }
   )
    lgbmr_test_metrics = test_eval_result.metrics

    loss = -lgbmr_val_metrics["val_r2_score"]

    # Truncate metric key names so they can be displayed together
    lgbmr_val_metrics = {k.replace("val_", ""): v for k, v in lgbmr_val_metrics.items()}
    lgbmr_test_metrics = {k.replace("test_", ""): v for k, v in lgbmr_test_metrics.items()}

    return {
      "loss": loss,
      "status": STATUS_OK,
      "val_metrics": lgbmr_val_metrics,
      "test_metrics": lgbmr_test_metrics,
      "model": model,
      "run": mlflow_run,
    }

Hyperparameter Tuning

This section configures the hyperparameter search space for optimizing the LightGBM Regressor using Hyperopt. The parameters specified in the space dictionary are constant values, but they can be modified to define a broader search space. For instance, discrete parameters like max_depth can be defined using hp.choice() to allow multiple possible values, enabling the optimization algorithm to explore a wider range of configurations. This flexibility is crucial for achieving the best model performance by systematically evaluating different hyperparameter combinations. Additionally, to improve optimization results, the max_evals parameter in the fmin function should be increased, allowing more trials during hyperparameter tuning:

space = {
  "colsample_bytree": 0.6291220708293621,
  "lambda_l1": 1.9107607188942655,
  "lambda_l2": 227.3523728055285,
  "learning_rate": 1.1282546701309717,
  "max_bin": 289,
  "max_depth": 3,
  "min_child_samples": 29,
  "n_estimators": 51,
  "num_leaves": 43,
  "subsample": 0.7122258812404411,
  "random_state": 990227116,
}

Running the Trails and Selecting the Best Model

The next section performs hyperparameter optimization using Hyperopt’s fmin() function, which searches for the best-performing model by evaluating different hyperparameter configurations. The Trials() object is used to keep track of all hyperparameter tuning runs, storing key details such as model performance metrics, selected hyperparameters, and MLflow run information. The optimization process applies the Tree-structured Parzen Estimator (tpe.suggest), a Bayesian optimization algorithm that efficiently explores the search space. The max_evals parameter, currently set to 1, determines the number of trials. Increasing this value allows for a broader and more thorough search, improving the likelihood of finding an optimal model.

Once the hyperparameter tuning process completes, the best trial is extracted from the trials object. The best model instance (model) and its corresponding MLflow experiment run (mlflow_run) are retrieved for further evaluation. The performance metrics for the validation and test sets are displayed using a Pandas DataFrame, providing a quick comparison of the model’s generalization capability. Finally, scikit-learn’s set_config(display="diagram") is used to visually represent the pipeline structure, allowing for an intuitive understanding of the transformations and model architecture. This step ensures that the best model from AutoML is selected and ready for deployment or further fine-tuning:

trials = Trials()
fmin(objective,
     space=space,
     algo=tpe.suggest,
     max_evals=1,  # Increase this when widening the hyperparameter search space.
     trials=trials)

best_result = trials.best_trial["result"]
model = best_result["model"]
mlflow_run = best_result["run"]

display(
  pd.DataFrame(
    [best_result["val_metrics"], best_result["test_metrics"]],
    index=["validation", "test"]))

set_config(display="diagram")
model

The final pipeline and model are shown in Figure 5-13.

Figure 5-13. Pipelines and Model for the California Housing Dataset AutoML run

Other Steps

The remainder of the notebook focuses on the ML lifecycle. Once the optimal LightGBM model is trained, it is logged to MLflow and can be registered in the MLflow Model Registry for version control and deployment. The notebook includes code snippets for registering the model, loading it for inference, and retrieving specific versions from the registry. Additionally, SHAP (Shapley Additive Explanations) is integrated to provide feature importance insights, though it is disabled by default and can be enabled to analyze the impact of each feature on predictions. Finally, the notebook demonstrates how to load the trained model from MLflow and perform inference on new data, ensuring seamless integration into production workflows. I encourage the readers to review these sections on their own time.

Summary

As machine learning becomes increasingly important across industries, Databricks AutoML provides a solution for automating complex model-building tasks, making them accessible to a wider audience. This chapter introduces AutoML’s capabilities in handling data preprocessing, feature selection, model training, hyperparameter tuning, and deployment with minimal manual intervention. By leveraging AutoML, users can build and optimize regression, classification, and time-series forecasting models efficiently. The chapter explores both the AutoML UI and Python API, demonstrating how they simplify the model development workflow. Additionally, it covers key AutoML functionalities, such as handling missing data, detecting class imbalances, integrating the Feature Store, and visualizing feature importance using SHAP values. These features allow users to rapidly develop machine learning models while ensuring explainability and reproducibility.

A significant portion of this chapter is dedicated to a real-world application of AutoML using the Kaggle Titanic dataset and the California Housing dataset. It provides hands-on guidance on launching AutoML experiments, configuring training settings, and interpreting results. The best AutoML-selected model, a LightGBM Regressor, is analyzed in detail, covering its preprocessing steps, hyperparameter tuning process, and model evaluation using MLflow. The chapter also explains how AutoML automates the train-validation-test split, registers the final model in the MLflow Model Registry, and provides tools for deploying and serving predictions. Readers will walk away with a comprehensive understanding of how Databricks AutoML optimizes machine learning workflows, allowing them to build, evaluate, and deploy models with minimal coding effort while maintaining transparency and control over the modeling process.

Chapter 6. The Feature Store

A Note for Early Release Readers

With Early Release ebooks, you get books in their earliest form—the author’s raw and unedited content as they write—so you can take advantage of these technologies long before the official release of these titles.

This will be the 6th chapter of the final book. Please note that the GitHub repo will be made active later on.

If you have comments about how we might improve the content and/or examples in this book, or if you notice missing material within this chapter, please reach out to the editor at jbleiel@oreilly.com.

The Databricks Feature Store serves as a critical component for feature engineering and management in the Lakehouse architecture. Feature engineering remains a cornerstone of machine learning success, as models depend on well-preprocessed features to deliver accurate predictions. The Feature Store provides a centralized repository for managing features, ensuring consistency between training and serving workflows, and enabling collaboration across teams. The Feature store integrates with Delta tables and Unity Catalog, ensuring governance, metadata tracking, and reusability of features.

In this chapter, you will explore how the Feature Store enhances multiple aspects of the ML lifecycle. This includes creating and managing feature tables, using feature lookups to generate training datasets, and enabling real-time inference with online feature stores. The integration of the Feature Store with MLflow ensures that feature metadata is retained throughout the deployment process, streamlining scoring and enabling reproducible workflows. You conclude this chapter with a practical example where we use the Lending Club dataset to create a Feature Table, and use the features in the table to create a model which predicts a loan default.

The Role of Feature Engineering in Machine Learning

Machine learning models do not operate on raw data. Instead, they depend on features (which are specific attributes or representations extracted from raw data) that encapsulate the underlying patterns and relationships needed for accurate predictions. Effective feature engineering involves:

Feature creation

Derive new variables that express the key behaviours of the problem domain (e.g., customer-lifetime value, 7-day rolling averages, or TF-IDF scores).

Feature extraction

Pull the most informative attributes from complex data (text, images, sensor traces) and encode them in a form a model can digest, trimming away incidental noise.

Feature selection

Keep what matters and drop the rest, pruning collinear or low-value columns so the model trains faster, generalises better, and stays interpretable.

In the Lakehouse architecture, these tasks can be performed seamlessly using Delta tables as a foundation for high-quality data storage and processing. Furthermore, feature engineering facilitates model explainability by ensuring features are interpretable and aligned with business objectives.

Feature engineering presents unique challenges that can hinder productivity and model performance if not managed properly:

Ensuring features used during training are identical during inference is critical to maintain consistency. This minimizes discrepancies and ensures that model performance during inference matches training expectations.

Collaboration is enhanced by enabling teams to reuse features across projects without duplicating efforts, fostering efficiency and teamwork.

Scalability is achieved by managing large-scale datasets while supporting both batch and real-time workflows, ensuring the system can handle diverse use cases.

Governance and compliance add another layer of complexity, requiring robust systems to track data lineage, metadata, and access permissions.

Delta Tables: The foundation of Feature Engineering

Delta tables in the Lakehouse serve as the backbone for feature engineering, providing a scalable, transactional storage layer that ensures data reliability and performance. With ACID transactions, Delta tables guarantee data consistency and reliability, making them a robust foundation for feature engineering. Additionally, the time travel functionality allows teams to retrieve historical versions of data, enabling reproducible ML experiments and ensuring traceability. Seamless integration with Spark facilitates efficient and scalable data processing for ML workflows, while schema enforcement ensures that feature tables conform to predefined structures, reducing errors during training and inference.

Feature Tables: The core of the Feature Store

Feature tables store records containing engineered features and are backed by Delta tables, enriched with metadata to enhance their utility. These tables enable efficient feature lookups, simplifying the process of retrieving features for both training and inference while reducing manual overhead and improving accuracy. Features can also be tagged and organized for easy discovery, fostering collaboration and cross-team usage. Lineage tracking maintains detailed records of data sources, transformations, and dependencies, ensuring transparency and auditability. Furthermore, automated monitoring ensures that feature tables are updated and consistent with the underlying data sources, reducing the need for manual intervention.

Integration with Unity Catalog

The Databricks Feature Store integrates seamlessly with Unity Catalog, elevating the governance and discoverability of features. Unity Catalog provides centralized governance by registering all feature tables under a unified namespace, enabling consistent access controls and comprehensive data lineage tracking across teams. Fine-grained access management allows administrators to define which users or groups can access specific feature tables, enhancing security and compliance. Features registered with Unity Catalog carry rich metadata, including details about source tables, transformations, and associated models, offering invaluable insights for audits and compliance. Unity Catalog also enables cross-workspace sharing of features, allowing teams across an organization to discover and leverage existing features seamlessly for their ML pipelines. Role-based access further ensures compliance with industry regulations while promoting ease of use.

By leveraging Unity Catalog, the Feature Store ensures that features are not only reusable but also compliant with organizational policies and industry regulations.

Feature Engineering in the Lakehouse

The Databricks Feature Store addresses these challenges head-on by providing a centralized platform for managing feature engineering processes. It is a purpose-built solution for feature engineering in a Lakehouse environment. By centralizing feature management, it empowers data scientists, ML engineers, and analysts to:

Maintain feature consistency across training and serving.

Reuse features across teams and projects.

Leverage the Lakehouse’s scalability for batch and real-time workflows.

Ensure governance and metadata tracking through its integration with Unity Catalog.

Figure 6-1 provides an overview of how the Databricks Feature Store integrates into the ML workflow:

Figure 6-1. Integration of the Databricks Feature Store in the ML workflow

Offline and Online Stores

The Feature Store supports both offline and online workflows, catering to a wide range of ML use cases:

The offline store is used for batch training and inference, with features materialized in Delta tables to ensure scalability and high performance for training datasets.

The online store supports real-time inference using low-latency databases such as Azure Cosmos DB. Features are published from the offline store to enable dynamic, real-time predictions for applications like fraud detection and personalized recommendations.

Feature Consistency and Reusability

The Databricks Feature Store ensures feature consistency by serving as a centralized repository for feature metadata, which is stored alongside models to streamline deployment and scoring processes. This approach allows teams to reuse features across multiple ML pipelines, reducing redundant efforts and significantly improving productivity. Moreover, the Feature Store supports hybrid workflows that seamlessly combine batch and real-time pipelines, enabling organizations to address a diverse range of business scenarios effectively and efficiently.

Machine Learning Workflow with the Feature Store

Feature engineering is an iterative and systematic process that integrates deeply into the ML development lifecycle. The steps below outline how the Databricks Feature Store is typically leveraged to streamline and enhance machine learning workflows.

Step 1: Creating and Managing Feature Tables

Features are generally created from curated Delta tables and then stored in the Feature Store to ensure consistency and reusability. Using the Feature Store API, teams can define and populate feature tables with preprocessed features, ensuring that they align with the specific requirements of machine learning models. As datasets evolve, new features can be added or existing ones merged to accommodate changing business needs. To maintain the reliability of feature tables, automated tools are employed to monitor feature health and ensure that feature tables remain synchronized with their source data, minimizing potential errors and inconsistencies.

Step 2: Training with Features from the Feature Store

Feature lookups streamline the creation of training datasets by automating the process of combining raw data with feature tables. This automation reduces manual effort and improves efficiency. Furthermore, the metadata associated with these features ensures reproducibility, allowing data scientists to track and reuse features across experiments and projects. Performance monitoring is also a critical aspect of this step, as it enables teams to evaluate how updates to features impact model accuracy and other performance metrics. This iterative feedback loop ensures that models are continuously optimized to deliver better results. Figure 6-2 illustrates the process of training models using an offline feature store.

Figure 6-2. Training models with an offline feature store

Step 3: Real-Time Inference

In real-time inference workflows, features are published from the offline store to an online store to support dynamic, low-latency predictions. This approach is particularly vital for applications such as fraud detection and personalized recommendations, where rapid response times are essential. The integration of MLflow simplifies the deployment and monitoring of models, ensuring that feature dependencies are managed efficiently throughout the ML lifecycle. Additionally, feature freshness is maintained through automated pipelines that continuously update the online store with the latest data, ensuring the accuracy and relevance of predictions. Figure 6-3 depicts real-time inferencing with an online feature store.

Figure 6-3. Real time inferencing with an online Feature Store

Example - Lending Club Feature Store

In this hands-on example, we’ll explore how the Feature Store integrates seamlessly within the lakehouse architecture to streamline feature engineering and model development. We’ll utilize the Lending Club Dataset, a rich source of financial data commonly used for credit risk modeling and loan default prediction.

First, you’ll extract meaningful features from the raw dataset and organize them into a structured Feature Store Table. This table will be governed and maintained by Unity Catalog, ensuring consistent data governance, lineage tracking, and access control across your ML workflows.

Next, we’ll demonstrate how to effortlessly consume features from the Feature Store to build a machine learning model. By leveraging this centralized feature repository, you’ll not only simplify model training but also ensure feature consistency between training and inference. Finally, we’ll evaluate the model’s performance, focusing on accuracy metrics, to illustrate the impact of well-managed features in driving reliable predictive outcomes.

Introduction to the Lending Club Dataset

The Lending Club dataset is a comprehensive collection of loan data issued by LendingClub, a prominent peer-to-peer lending platform. This dataset is frequently utilized in financial modeling and machine learning applications due to its richness in features and real-world relevance. It encompasses detailed records of loan applications, borrower profiles, and loan performance metrics, making it an ideal candidate for demonstrating feature engineering and management in a Lakehouse environment.

Key Attributes of the Lending Club Dataset:

Loan Details:

Information such as loan amount (loan_amnt), interest rate (int_rate), loan term, and loan status (loan_status), indicating whether a loan is fully paid, charged off, or current.

Borrower Information

Demographic and financial attributes like annual income (annual_inc), debt-to-income ratio (dti), credit score ranges, and employment history.

Credit and Payment History

Features including the number of open accounts (open_acc), revolving credit utilization (revol_util), and delinquencies (delinq_2yrs).

Derived Financial Metrics

Indicators like the ratio of income to loan amount, payment schedules, and risk scores derived from the available data.

In this section, you will use the Lending Club dataset to illustrate how to create, manage, and serve features using the Databricks Feature Store. This hands-on example will guide you through the process of feature engineering, storing features for reuse, and integrating them into machine learning workflows within a Lakehouse environment.

Approach

In this example, you will take a structured, end-to-end approach to demonstrate how the Feature Store enhances the machine learning lifecycle within a lakehouse environment. The workflow follows these six key steps:

Load and Preprocess the Lending Club Dataset. You will begin by loading the Lending Club Dataset into your environment. This step involves cleaning and preprocessing the data, including handling missing values, encoding categorical variables, and normalizing numerical features to ensure it’s ready for feature extraction.

Perform Feature Engineering. Next, you will derive new features from the raw dataset. This step focuses on creating meaningful, domain-specific features that can improve the model’s predictive performance. Feature engineering may include transformations, aggregations, or the creation of interaction terms.

Prepare the Features for Model Generation. Once engineered, the features will be curated and stored in a Feature Store Table managed by Unity Catalog. This ensures that the features are versioned, governed, and easily accessible for both training and inference.

Build the Model Using Features from the Feature Store. You will build a machine learning model by directly consuming features from the Feature Store Table. This demonstrates the efficiency of using a centralized feature repository, ensuring that the same features are consistently applied across different stages of the ML pipeline.

Create Visualizations for Model Evaluation. To interpret the model’s performance, you will generate visualizations that highlight key metrics and insights. This could include confusion matrices, ROC curves, or feature important plots, providing a clear understanding of how well the model is performing.

Log the Model in the Feature Store. Finally, you will log the trained model within the Feature Store. This step links the model to the specific features used during training, ensuring traceability and simplifying future deployment or retraining efforts.

Step 1 - Loading and preprocessing the Lending Club Dataset

As a first step you will load the data. The Lending Club data is available as a public dataset from Databricks. Since the dataset is available in parquet format, we can use the spark.read.parquet() spark function as is shown below:

# Load and Preprocess the Dataset with feature engineering
lending_club_path = "/databricks-datasets/samples/lending_club/parquet"
df = spark.read.parquet(lending_club_path)

A key requirement for creating a Feature Store Table is the presence of a unique identifier (primary key) for each record. This identifier ensures that features can be accurately referenced, tracked, and joined during both training and inference, maintaining consistency across the machine learning pipeline.

The Lending Club Dataset does not include a pre-existing unique identifier, which poses a challenge when preparing data for the Feature Store. To address this, we’ll generate a unique ID using Spark’s monotonically_increasing_id() function. This method creates a unique, non-sequential identifier for each row in the dataset, which will serve as the primary key in our Feature Store Table.

Here’s how you can generate and assign the unique ID:

# Generate a globally unique ID for each row
df = df.withColumn("unique_id", monotonically_increasing_id())

In many real-world datasets such as the lending club, numerical values such as interest rates or utilization ratios are often stored as strings with symbols like the percentage sign (%). To perform meaningful numerical analysis or feed these values into machine learning models, it’s essential to clean and convert them into appropriate numeric data types.

In the Lending Club Dataset, two columns, int_rate (interest rate) and revol_util (revolving credit utilization are stored as percentage strings (e.g., ’13.56%'). To convert these columns into a numerical format, you follow a multi-step cleaning process using Apache Spark functions. You first use trim to remove any whitespace, we then use a regular expression to remove the percentage signs (‘%’), and finally we cast the remaining string to a double.

df = df.withColumn("int_rate", regexp_replace(trim(col("int_rate")), "%", "").cast("double"))
df = df.withColumn("revol_util", regexp_replace(trim(col("revol_util")), "%", "").cast("double"))

You now have a Spark dataframe with clean data and a unique id, so you are now ready to turn this dataframe into a reusable set of features in a feature table.

Step 2 - Feature Engineering - Building you Feature Table

One of the key advantages of using a Feature Store within a lakehouse architecture is its ability to centralize, manage, and reuse features across various machine learning workflows. In this section, we’ll demonstrate how to build and incrementally expand a Feature Store Table using the Lending Club Dataset. This process will illustrate not only how to create an initial set of features but also how to seamlessly merge additional features into the same table, ensuring consistent feature management over time.

Two-Pass Approach to Feature Table Creation

To simulate a real-world machine learning pipeline where features are often developed in stages, you will adopt a two-pass approach:

Initial Feature Selection and Table Creation. In the first pass, you will select a subset of key features that are immediately relevant to our credit risk modeling task. These features will be written to a new Feature Store Table, with the unique_id column serving as the primary key to uniquely identify each loan record.

Merging Additional Features into the Feature Table. In the second pass, you will engineer additional features that provide deeper insights into the borrower’s credit profile. These new features will be merged into the existing Feature Store Table using the primary key to ensure data integrity. This demonstrates how the Feature Store allows for incremental feature additions without disrupting previously stored data.

This step-by-step process highlights how a feature store facilitates collaborative feature engineering, enabling teams to iteratively enhance datasets while maintaining governance, traceability, and consistency.

Phase 1 - Selecting and Writing Initial Features

We begin by selecting an initial set of features from the Lending Club Dataset that are crucial for credit risk assessment. These include loan-specific attributes like the loan amount and interest rate, as well as borrower-related financial indicators such as annual income and debt-to-income ratio (DTI):

# Feature Engineering
# Selecting relevant features for the model.
# The 'unique_id' column serves as the primary key for uniquely
# identifying each row in the Feature Store.
initial_features_df = df.select(
    col("unique_id"),   # Unique identifier for each loan record
    col("loan_amnt"),   # The amount of money requested by the borrower
    col("int_rate"),    # Interest rate of the loan, converted to numeric format
    col("annual_inc"),  # The Borrower's annual income
    col("dti"),         # Debt-to-income ratio, a key metrics for credit risk
    col("revol_util"),  # Revolving credit line utilization rate
    col("open_acc")     # Number of open credit lines in the borrow's credit file
)

Phase 2 - Setting up Unity Catalog for Feature Storage

Before writing features to the Feature Store, you need to ensure that our catalog and schema exist in Unity Catalog, which governs feature access and lineage across the workspace. This ensures that our feature table is created in a properly managed environment, maintaining compliance with organizational governance standards.

%sql
CREATE SCHEMA IF NOT EXISTS book_ai_ml_lakehouse.feature_store_db;

Next, you define a fully qualified name for our Feature Store Table following the format catalog.schema.table_name to maintain clarity and consistency in data organization.

# We will be using Unity Catalog to store our Feature Table.
FEATURE_TABLE_NAME = "lending_club_loan"
feature_table_name = f"{CATALOG_NAME}.{FEATURE_STORE_DB}.{FEATURE_TABLE_NAME}"
print(f"Our feature table will be named: `{feature_table_name}`")

Output:

Our feature table will be named: `book_ai_ml_lakehouse.feature_store_db.lending_club_loan`

Step 3- Feature Engineering

With the initial features selected and the schema in place, we now create the Feature Store Table using Databricks Feature Store APIs.

Step 1 - Create the feature table

The unique_id column is designated as the primary key to uniquely identify each record, and a description is provided for better traceability.

# Write the initial subset of features to the Databricks Feature Store
fs.create_table(
    name=feature_table_name,
    primary_keys=["unique_id"],
    df=initial_features_df,
    description="Initial subset of engineered features for Lending Club Loan dataset"
)

Output:

<FeatureTable: name='book_ai_ml_lakehouse.feature_store_db.lending_club_loan', table_id='b93377b9-0ec3-4e52-9f21-506fb9abc2a2', description='Initial subset of engineered features for Lending Club Loan dataset', primary_keys=['unique_id'], partition_columns=[], features=['unique_id',
 'loan_amnt',
 'int_rate',
 'annual_inc',
 'dti',
 'revol_util',
 'open_acc',
 'delinq_2yrs',
 'pub_rec',
 'total_acc',
 'mort_acc'], creation_timestamp=1739229766072, online_stores=[], notebook_producers=[], job_producers=[], table_data_sources=[], path_data_sources=[], custom_data_sources=[], timestamp_keys=[], tags={}>

At this point, you can navigate to the Databricks Console and view your Feature Table under the Feature Store section. As shown in Figure 6-4, the table lending_club_loan will be listed with the initial set of features.

Figure 6-4. The Feature Table listed in the Feature store

When you click on the lending_club_loan table, you will see the initial set of features listed,as shown in Figure 6-5.

Figure 6-5. The details of the lending_club_loan feature table.

Step 2 - Merging Additional Features into the Feature Store

In real-world scenarios, feature engineering is often iterative, with new features being added as more data becomes available or new insights are uncovered. Here, you demonstrate how to append additional features such as delinquency history and total number of credit accounts to the existing Feature Store Table.

# Additional feature engineering
# Selecting additional relevant features to enhance model performance.
# These features provide deeper insights into the borrower's credit history and financial health.
# Loading additional features into a PySpark DataFrame
additional_features_df = spark.createDataFrame(
    df.select(
        col("unique_id"),
        col("delinq_2yrs"),
        col("pub_rec"),
        col("total_acc"),
        col("mort_acc")
    ).rdd
)

To integrate these features into the existing Feature Store Table, you use the merge mode. This operation ensures that the new features are appended based on the unique_id, while preserving existing data and maintaining data lineage.

# Appending additional features to the Feature Store
# This step integrates the newly engineered features into
# the existing feature table in the Feature Store.
fs.write_table(
    name=feature_table_name,
    df=additional_features_df,
    mode="merge"
)

By leveraging the merge functionality, you ensure that your Feature Store remains a dynamic, evolving resource that supports ongoing model development while maintaining historical data integrity.

Step 3 - Verifying the Feature Store Update

After merging the new features, refresh the Feature Store page in the Databricks Console. You’ll now see the updated lending_club_loan Feature Table with the additional features appended, as shown in Figure 6-6. This demonstrates how easily features can be added to existing tables, making the Feature Store a powerful tool for collaborative and iterative ML development.

Figure 6-6. The Feature Table with the additional features added

Step 4 - Prepare the Features for Model Creation

Once features have been engineered and stored in the Feature Store, the next crucial step is to prepare them for machine learning model development. This involves reading the features from the Feature Store, aligning them with the target variable (in this case, loan default status), and ensuring the data is clean and ready for model training. Proper feature preparation is essential to ensure model accuracy, data integrity, and reproducibility.

In this section, we’ll walk through the process of preparing features from the Lending Club Dataset for a loan default prediction model. The workflow includes the following steps:

Reading Features from the Feature Store

Defining the Target Variable (Loan Default Status)

Merging Features with the Target Variable

Cleaning and Preparing Data for Model Training

Handling Missing Values and Ensuring Numeric Data

This structured approach ensures that the dataset is in optimal shape for training a reliable and robust machine learning model.

Step 5 - Reading Features from the Feature Store

To begin, you will retrieve the engineered features from the Feature Store using the read_table API provided by Databricks Feature Store. The features are then converted into a Pandas DataFrame to facilitate further local processing and compatibility with common machine learning libraries such as scikit-learn.

# Reading features from the Feature Store
# Features are read from the Feature Store and converted into a Pandas DataFrame for local processing.
features = fs.read_table(feature_table_name).toPandas()

By leveraging the Feature Store, we ensure that the features are consistent with what will be used during model inference, reducing the risk of training-serving skew.

Step 6 - Defining the Target Variable

For this modeling task, we aim to predict whether a loan will default. In the Lending Club Dataset, a loan is considered defaulted if its loan_status is marked as “Charged Off”. We create a new binary column, defaulted, where:

True indicates a defaulted loan (i.e., loan_status == “Charged Off”).

False indicates a non-defaulted loan.

# Load target variable (e.g., predicting loan default)
# The target variable 'defaulted' is created based on the 'loan_status' column.
# If the loan status is 'Charged Off', it indicates a default, and the value
# is set to True; otherwise, it is False.
# The 'unique_id' ensures the target variable aligns with the features
# in the Feature Store for proper model training and evaluation.
# The data is then converted to a Pandas DataFrame for compatibility with scikit-learn.
labels = df.select("unique_id", (col("loan_status") == "Charged Off").alias("defaulted"))

The unique_id ensures that each target label is correctly aligned with its corresponding feature set during the merging process.

Step 7 - Merging the Features with the target variable

To create a complete dataset for model training, you merge the features retrieved from the Feature Store with the defaulted target variable. The merge is performed using the unique_id to ensure accurate alignment between the features and their respective labels.

# Merge features with labels
# Combines the engineered features from the Feature Store
# with the target variable 'defaulted'.
# The merge is performed using 'unique_id' to ensure each feature set
# is accurately linked to its corresponding target label.
data = pd.merge(features, labels.toPandas(), on="unique_id")

This combined dataset now contains all the information needed to train and evaluate a machine learning model.

Step 8 - Handling missing Target Values

Before proceeding to model training, it’s important to handle any missing target values. Since machine learning algorithms cannot train on records with undefined target labels, you will drop any rows where the defaulted status is missing.

# Drop rows with missing target values
# Ensures that any records without a defined 'defaulted' status are removed.
# This step is crucial because machine learning models cannot train on missing target labels.
data = data.dropna(subset=["defaulted"])

This ensures that our model is trained on complete, reliable data.

Step 9 - Separating the Features from the Target Values

With the dataset cleaned, we now separate the features from the target variable. We’ll drop the unique_id column (as it doesn’t contribute to the prediction) and the defaulted column (which will serve as the target variable during training).

# Prepare features for model training
# Drops 'unique_id' (as it's not useful for prediction) and 'defaulted' (the target variable).
# This ensures the model only trains on relevant feature columns.
X = data.drop(columns=["unique_id", "defaulted"])

Step 10 - Converting all Features to Numeric Types

Most machine learning algorithms require numerical input. To ensure compatibility, you convert all feature columns to numeric types. Any non-convertible values are coerced to NaN, which will be handled in the next step.

# Ensure all features are numeric
# Converts all feature columns to numeric types, coercing errors to NaN if necessary.
# This is essential because machine learning algorithms require numeric input.
X = X.apply(pd.to_numeric, errors='coerce')

Step 11 - Imputing Missing Values

It’s common for real-world datasets to contain missing values. To handle these, we use a SimpleImputer with a mean strategy, which replaces missing values with the mean of the respective columns. This ensures the dataset is complete and ready for model training.

# Impute missing values using mean strategy
# The SimpleImputer replaces missing values (NaNs) in the feature set with the mean of each column.
# This is a common preprocessing step to handle incomplete data, ensuring that machine learning algorithms
# can be trained without errors due to missing values.
imputer = SimpleImputer(strategy='mean')
X_imputed = imputer.fit_transform(X)

By imputing missing values, we avoid potential errors during model training and ensure that the model can generalize better to new data.

Step 12 - Defining the Target Variable

With the dataset fully prepared, the next step is to isolate the target variable for model training. In this case, the target variable is defaulted, which indicates whether a loan has been charged off (i.e., defaulted) or not. This binary classification will allow us to build a predictive model that distinguishes between loans likely to default and those that are not.

The following line of code extracts the defaulted column from the dataset and assigns it to the variable y, which is the conventional notation for the target variable in machine learning workflows:

# Set the Target variable
y = data["defaulted"]

With the features and target variable prepared, the next step is to build and train a machine learning model to predict loan defaults.

Step 4 - Build the Model with the Features of the Feature Store

With our features engineered, stored, and prepared, it’s time to transition from data preprocessing to model building. In this section, you will construct a machine learning model to predict loan defaults using the features retrieved from the Feature Store. This process involves splitting the data, selecting an appropriate model, training it, and evaluating its performance. Each step is crucial to ensure that the model not only fits the data well but also generalizes effectively to unseen cases.

Step 1 - Splitting the Dataset for Training and Testing

Before training a machine learning model, it’s essential to split the dataset into training and testing sets. This ensures that the model is evaluated on data it hasn’t seen during training, providing an unbiased assessment of its predictive performance. We’ll use the widely trusted train_test_split function from scikit-learn to perform this s

# Split the dataset
# Divides the dataset into training and testing sets to evaluate model performance.
# 80% of the data is used for training, and 20% is reserved for testing.
# The 'random_state' ensures reproducibility by setting a seed for the random number generator.
X_train, X_test, y_train, y_test = train_test_split(X_imputed, y, test_size=0.2, random_state=42)

Step 2 - Choosing the Random Forest Classifier

For this project, we have selected the Random Forest Classifier, an ensemble learning method renowned for its robustness and accuracy, particularly in handling complex datasets like those found in financial services. So, why did we select the Random Forest for this type of Modeling?

We select the Random Forest for this model. The main driver for this is the fact that the Lending Club Dataset has some unique characteristics:

High Dimensionality

It includes a wide variety of features such as loan amount, interest rates, credit history, and income.

Noisy and Incomplete Data

Financial datasets often contain missing values, outliers, and noisy entries.

Non-Linear Relationships

Borrower characteristics and loan outcomes usually exhibit non-linear interactions.

Imbalanced Classes

Loan default datasets often have far fewer defaults than fully paid loans.

The Random Forest Classifier addresses these challenges effectively:

By selecting random subsets of features at each split, Random Forest reduces overfitting and improves generalization.

Robust to Noisy Data: The ensemble nature of Random Forest mitigates the influence of outliers and missing values through surrogate splits.

Captures Complex Patterns: It models non-linear relationships without extensive feature engineering.

Feature Importance Insights: Random Forest provides feature importance scores, offering valuable interpretability, crucial for understanding risk factors in credit modeling.

Adaptable to Imbalanced Data: Techniques like class weighting or resampling can be integrated into Random Forest to address class imbalances.

Moreover, while Random Forest can be computationally intensive, its parallelizable nature makes it scalable for large datasets, particularly when deployed on cloud platforms like Databricks.

Step 3 - Training the Random Forest Model

With the training data prepared and the model selected, we can now train the Random Forest Classifier. We’ll initialize the model with 100 decision trees (estimators), striking a balance between performance and computational efficiency.

# Train a Random Forest classifier
# Initializes the Random Forest model with 100 decision trees (estimators).
# The 'random_state' parameter ensures reproducibility by fixing the random number generation.
# The model is then trained using the training feature set (X_train) and target labels (y_train).
# The target variable is explicitly converted to integers to ensure
# compatibility with the classifier.
model = RandomForestClassifier(n_estimators=100, random_state=42)
model.fit(X_train, y_train.astype('int'))

In this code:

n_estimators=100 specifies the number of decision trees in the forest, providing a robust ensemble that reduces variance.

random_state=42 ensures that the random processes involved in the model (like feature selection at each split) are reproducible.

y_train.astype('int') ensures compatibility with the classifier, converting boolean labels into integers (0 for non-default and 1 for default).

This produces the output as shown in Figure 6-7.

Figure 6-7. Model creation output

After training, the model is ready to make predictions and assess its performance.

Step 4 - Making Predictions and Evaluation Model Performance

With the model trained, the next step is to evaluate its performance on the test dataset. This involves predicting loan defaults and calculating the accuracy score, a metric that measures the proportion of correct predictions.

Before making predictions, you ensure the target labels in the test set are in the correct numerical format:

# Predict and evaluate
# Ensure target labels are binary and consistent
# Converts the target labels in y_test to integers if they are in boolean format.
# This step is crucial because scikit-learn's classification metrics expect numerical labels (0 and 1).
y_test = y_test.astype(int)

Now, you use the trained model to predict loan default outcomes:

# Make predictions on the test dataset
# Uses the trained Random Forest model to predict loan default outcomes on the test set.
y_pred = model.predict(X_test)

To ensure data integrity, we verify that both the actual and predicted values are binary (0 for non-default and 1 for default):

# Check for unexpected values
# Prints the unique values in y_test and y_pred to ensure both contain only binary values (0 and 1).
# This step helps identify any anomalies or data inconsistencies that could affect evaluation.
print("Unique values in y_test:", y_test.unique())
print("Unique values in y_pred:", pd.Series(y_pred).unique())

You calculate the accuracy of the model, which represents the proportion of correct predictions:

# Evaluate model accuracy
# Calculates the accuracy of the model by comparing predicted
# values (y_pred) with actual values (y_test).
# Accuracy represents the proportion of correct predictions out of all predictions made.
accuracy = accuracy_score(y_test, y_pred)
print(f"Model Accuracy: {accuracy}")

This produces the following output:

Unique values in y_pred: [0 1]
Model Accuracy: 0.9077280811585203

The model achieves an impressive accuracy of 90.77%, indicating strong performance in predicting loan defaults. However, while accuracy is a useful initial metric, it doesn’t provide the complete picture, especially in imbalanced datasets where one class (e.g., non-defaults) significantly outweighs the other. To gain deeper insights, we’ll explore additional evaluation metrics like the confusion matrix, precision, recall, and F1 score in the next section on model visualizations.

Step 5 - Inferring the Model Signature with MLflow

When deploying machine learning models, especially in production environments, it’s essential to define a model signature. A model signature describes the input schema (features the model expects) and output schema (predictions or outputs the model produces). This ensures consistency between training and inference, reducing the risk of errors due to mismatched data types or structures.

In this example, we use MLflow’s infer_signature function to automatically capture the structure of both the input features and the model’s predictions:

from mlflow.models import infer_signature

# Infer model signature for Unity Catalog
signature = infer_signature(X_test, y_pred)

By capturing the model signature, we enhance model reproducibility, traceability, and deployment readiness, aligning with best practices for managing machine learning models in a lakehouse architecture.

Step 5 - Model Visualizations

While numerical metrics like accuracy provide a quick snapshot of model performance, visualizations offer deeper insights into both how the model makes predictions and how well it performs. Visual tools such as feature importance plots and confusion matrices not only enhance interpretability but also help in diagnosing issues like overfitting, class imbalance, or irrelevant features. In this section, we’ll focus on the key visualizations that illuminate the performance and decision-making process of our Random Forest Classifier.

The Feature Importance Plot

A Feature Importance Plot illustrates the relative importance of each feature in the model’s decision-making process. In ensemble methods like Random Forest, feature importance is determined based on how much each feature reduces impurity across the decision trees. Features that contribute more to splitting data and improving prediction accuracy are ranked higher.

A Feature Importance Plot provides insights into the following areas:

Model Interpretability

It helps explain which features most influence the model’s predictions, an essential aspect in domains like finance where interpretability is critical.

Feature Selection

Identifying unimportant features can guide you in simplifying the model, reducing overfitting, and improving performance.

Business Patterns

Feature importance can reveal underlying patterns in the data that are valuable for business decisions—in this case, understanding what factors most contribute to loan defaults.

The Feature Importance Plot for our model is shown in Figure 6-8.

Figure 6-8. Feature Importance plot.

From the plot, we immediately notice that the most influential features in predicting loan defaults are:

dti (Debt-to-Income Ratio)

This measures the applicant’s debt burden relative to their income, a critical factor in assessing credit risk.

revol_util (Revolving Credit Utilization)

Indicates how much of their available revolving credit (like credit cards) the applicant is using, which can reflect financial stress.

total_acc (Total Number of Accounts)

Provides an overview of the applicant’s credit history, where a higher number of accounts may indicate greater financial complexity or risk.

These features not only drive the model’s predictions but also offer actionable insights into the key factors influencing loan repayment behavior.

The Confusion Matrix

While overall accuracy gives us a broad sense of model performance, it doesn’t tell the full story—especially in scenarios where the data is imbalanced, such as predicting loan defaults where non-defaulted loans far outnumber defaults. To gain a clearer understanding, we turn to the Confusion Matrix.

A Confusion Matrix is a performance evaluation tool that compares the model’s predicted classifications with the actual outcomes. It breaks down predictions into four categories:

True Positives (TP)

Correctly predicted defaults.

True Negatives (TN)

Correctly predicted non-defaults.

False Positives (FP)

Non-defaults incorrectly predicted as defaults (Type I error).

False Negatives (FN)

Defaults incorrectly predicted as non-defaults (Type II error).

The confusion matrix for your model is shown in Figure 6-9.

Figure 6-9. [Caption to come]

From the lower matrix, we can extract key performance metrics:

Accuracy

Measures the proportion of correct predictions out of all predictions made.

Precision

Indicates the percentage of loans predicted as defaults that were actually defaults.

Recall (Sensitivity)

Shows the proportion of actual defaults that were correctly identified by the model.

F1 Score

Provides a balance between precision and recall, especially useful when dealing with imbalanced datasets.

Model Performance Overview

Based on the visualizations, we can see that our model performs well, with an accuracy of over 90% and strong scores across precision, recall, and F1 metrics. The Feature Importance Plot highlights critical risk factors like debt-to-income ratio and revolving credit utilization, while the Confusion Matrix confirms the model’s reliability in distinguishing between defaulted and non-defaulted loans.

In the next section, you will log the model to the feature store, connecting the model to its source features

Step 6 - Log the Model in the Feature Store

Once the machine learning model has been trained and evaluated, the next critical step is to log and register the model for future use. This ensures the model is governed, versioned, and reproducible, making it easier to deploy, monitor, and manage across different environments. In this section, we’ll demonstrate how to log the trained Random Forest Classifier into MLflow and register it within Databricks Unity Catalog, linking it to the features stored in the Feature Store. This process guarantees that the model remains consistent with the data it was trained on, reducing the risk of training-serving skew.

Step 1 - Building the Unity Model Name

Before logging the model, we need to define its fully qualified name in Unity Catalog. This naming convention ensures that the model is uniquely identifiable and easily accessible across different Databricks workspaces.

# The unique identifier for the MLflow run where the model was logged.
# MLflow assigns a unique 'run_id' to each experiment run. This ID is critical
# for tracking the specific model version, its parameters, metrics, and artifacts.
run_id = "3cb01eb8044e47feb5d9dcaedbfc9fcd"

# The name assigned to the model when registering it in MLflow.
# This 'model_name' serves as a reference for the registered model in both MLflow
# and Unity Catalog, enabling consistent access across different environments.
model_name = "lending_club_default_model"

# The catalog name in Unity Catalog, representing the top-level namespace.
# Unity Catalog uses 'catalog_name' as the primary organizational layer
# for managing data and models across the lakehouse.
catalog_name = CATALOG_NAME  # Ensure CATALOG_NAME is defined earlier in your code.

# The schema name within the catalog, used for organizing registered models.
# The 'schema_name' helps categorize models within a specific catalog,
# promoting better governance and discoverability.
schema_name = FEATURE_STORE_DB  # Ensure FEATURE_STORE_DB is defined earlier in your code.

# Construct the fully qualified Unity Catalog model name.
# The format 'catalog.schema.model_name' ensures that the model is uniquely
# identifiable and governed across workspaces in Databricks.
unity_model_name = f"{catalog_name}.{schema_name}.{model_name}"

# Display the full model path in Unity Catalog for verification.
# This helps confirm that the model has been correctly registered
# and can be easily referenced for deployment or inference.
print(f"Unity Model Name: {unity_model_name}")

Output:

Unity Model Name: book_ai_ml_lakehouse.feature_store_db.lending_club_default_model

Step 2 - Setting the Unity Catalog URI

By default, MLflow uses its internal registry for tracking models. To leverage the advanced governance and cross-workspace capabilities provided by Unity Catalog, we need to set the registry URI to point to Databricks Unity Catalog:

# Set the registry URI to Databricks Unity Catalog.
# By default, MLflow uses its internal model registry for tracking models.
# This command overrides the default and instructs MLflow to use Databricks Unity Catalog
# as the backend for model management.
#
# Unity Catalog provides centralized governance, lineage tracking, and access control
# for models across all Databricks workspaces. By integrating MLflow with Unity Catalog,
# models benefit from enhanced security, consistent data governance, and
# cross-workspace discoverability.
mlflow.set_registry_uri("databricks-uc")

This ensures that all models logged from this point forward are stored and managed within Unity Catalog, benefiting from enterprise-grade security and governance features.

Step 3 - Logging the Model to the Feature Store

Now that the environment is configured, you proceed to log the trained model using MLflow and link it to the Feature Store. This step ensures that the model remains tightly coupled with the features it was trained on, enabling consistent and reproducible predictions during inference.

from mlflow import MlflowClient

# Start an MLflow run to log model artifacts and metrics
# This context manager initializes a new MLflow run, allowing you to log models, metrics,
# and other artifacts. The run provides a unique run ID that helps track and reference
# the model within the MLflow and Unity Catalog ecosystems.
with mlflow.start_run() as run:

    # Log the trained Random Forest model along with its signature
    # The model is logged using MLflow's sklearn integration. The 'signature' ensures
    # that the input-output schema of the model is recorded, maintaining consistency
    # between training and inference.
    mlflow.sklearn.log_model(clf, "model", signature=signature)

    # Log the model's accuracy metric
    # This captures the model's accuracy as a performance metric in MLflow,
    # allowing easy tracking, comparison, and evaluation across different experiments.
    mlflow.log_metric("accuracy", accuracy)

    # Register the model in Unity Catalog
    # This step registers the logged model in Unity Catalog for centralized governance
    # and cross-workspace accessibility. The 'run.info.run_id' links this registration
    # to the specific MLflow run, ensuring full traceability.
    model_version = mlflow.register_model(f"runs:/{run.info.run_id}/model", unity_model_name)

    # Define feature lookups for model deployment
    # The FeatureLookup specifies which features the model requires at inference time.
    # It maps the 'unique_id' key to features stored in the Feature Store, ensuring
    # consistency between training and serving.
    feature_lookups = [
        FeatureLookup(
            table_name=feature_table_name,  # Name of the feature table in Unity Catalog
            feature_names=["loan_amnt", "int_rate", "annual_inc", "dti", "revol_util", "open_acc"],  # Features used in the model
            lookup_key="unique_id"  # Primary key used for joining features during inference
        )
    ]

    # Link the model to the Feature Store for feature consistency during inference
    # This step logs the model in both MLflow and the Feature Store, associating it
    # with the features used during training. The 'training_set' ensures that the model
    # remains consistent with the feature data, enabling reproducibility and
    # simplifying deployment workflows.
    fs.log_model(
        model=clf,  # The trained Random Forest model
        artifact_path="model",  # Path within the MLflow run where the model is stored
        flavor=mlflow.sklearn,  # Specifies the MLflow flavor for scikit-learn models
        training_set=fs.create_training_set(
            df=df.select("unique_id", (col("loan_status") == "Charged Off").alias("defaulted")),  # Dataset with labels
            feature_lookups=feature_lookups,  # Defined feature lookups for inference
            label="defaulted"  # Target variable used during training
        ),
        registered_model_name=unity_model_name  # Fully qualified model name in Unity Catalog
    )

# Confirm successful model training, logging, and registration
print("Model training and logging to Unity Catalog completed.")

Step 4 - Viewing the Feature Store and the Model in the Databrick Console

After successfully logging the model and registering it within Unity Catalog, it’s important to verify that both the Feature Table and the model are correctly stored and governed in the Databricks environment. The Databricks Console offers an intuitive interface to navigate and manage data assets, including feature tables and machine learning models. This step ensures that all components are organized, versioned, and accessible for future reference, deployment, or collaboration across teams.

Navigating to the Feature Store and Model in Unity Catalog

To navigate to the feature store and model in Unity Catalog, you follow these steps as illustrated in Figure 6-10:

Access Unity Catalog:

From the Databricks workspace, navigate to the left-hand navigation pane.

Click on Catalog to open Unity Catalog, which serves as the central location for managing datasets, feature tables, and machine learning models.

Expand the Feature Store Schema:

In the catalog list, locate and expand the book_ai_ml_lakehouse catalog (or the catalog you’ve assigned).

Under the catalog, expand the feature_store_db schema. This schema contains both the Feature Table and the registered model associated with the Lending Club Dataset.

View the Feature Table and Model as shown in Figure 6-10.

In feature_store_db, you’ll see:

The Feature Table containing the engineered features for loan default prediction.

The registered model is named lending_club_default_model, which has been logged and linked to the features in the Feature Store.

Figure 6-10. The model and feature table in Unity Catalog.

In this figure, you can see both the lending_club_loan Feature Table and the lending_club_default_model under the feature_store_db schema in Unity Catalog.

Inspecting the Model in Unity Catalog

To inspect the model in the Unity Catalog, follow these steps:

Click on the Model:

Select lending_club_default_model from the list of models in the feature_store_db schema.

View Model Versions:

Unity Catalog supports versioning for models. Click on the model name to view its details.

You’ll see a list of all registered versions of the model. Select the most recent version to explore the specific model run.

Model Metadata and Artifacts:

After selecting the latest version, Unity Catalog provides detailed information about the model:

Model Artifacts: View the saved model files, such as the serialized model object and input-output signature.

Parameters and Metrics: Review the parameters used during training and performance metrics like accuracy.

Feature Store Integration: Confirm that the model is linked to the correct Feature Table, ensuring consistent feature retrieval during inference.

Lineage and Governance: Unity Catalog also displays the lineage of the model, showing which datasets and features contributed to its training, providing transparency and traceability.

This is illustrated in Figure 6-11.

Figure 6-11. The lending_club_default_model in Unity Catalog.

In this figure, the model details are displayed, showing key metrics, versioning information, and the linkage to the lending_club_loan Feature Table. This ensures full transparency and governance over the model lifecycle.

Summary

Feature engineering in the Lakehouse enables consistent and scalable machine learning workflows. The Databricks Feature Store, integrated with Unity Catalog, helps teams manage features efficiently across batch and real-time pipelines while ensuring consistency and collaboration. It supports governance, metadata management, and cross-workspace sharing, allowing organizations to maintain compliance and streamline operations.

About the Author

Bennie Haelen is a principal architect with Insight, a Microsoft and Databricks partner. As principal architect, his primary focus areas are modern data warehousing, machine learning, and generative AI on various commercial cloud platforms. Over his career, Bennie has overseen many data and AI projects in different application domains, such as healthcare, energy, and financial applications. He has also architected and delivered generative AI solutions that leverage both commercial and open source models.