This tutorial at PyData Global 2024, led by Nathan Kolbert, Senior Manager of ML Architecture at Peacock, guides participants through building a core machine learning platform from scratch using Python. The session emphasizes test-driven development (TDD) and domain-driven design (DDD) principles to create a scalable and testable system focused on personalized predictions for users. Key components explored include an inference engine, model repository, feature repository, and the concept of a message bus for decoupled service communication via event-driven design. The tutorial demonstrates how abstractions like the repository pattern and message handlers facilitate testing system logic independently of external infrastructure, even addressing complex end-to-end scenarios involving missing features or models. While not covering infrastructure, CI/CD, or advanced model training in depth, the session lays a foundational architectural blueprint focused on creating tightly bounded, collaborative services.
Building a Core ML Platform: From Features to Inference
0:00 Introduction: Nathan Kolbert, Senior Manager of ML Architecture at Peacock, introduces the tutorial at PyData Global on building a core ML platform from scratch.
0:30 Speaker Background: Nathan shares his background, transitioning from a musician (jazz drums) to an ML professional, with experience in data pipelines, feature stores, model deployment, and currently working on a global ML platform at Peacock.
1:30 Tutorial Goal: The aim is to build a core ML platform, specifically a monolithic yet modular system, focusing on the "from inference to features" approach using a test-driven method guided by domain-driven design.
3:17 ML Platform Components: The core components of an ML platform are defined: inference engine (predictions), feature store (feature storage), model registry (model storage/versioning), message bus (service communication), and potentially a model trainer.
4:56 Building Intuition: The focus is on building intuition for creating scalable, testable platforms, rather than directly covering infrastructure, CI/CD, cost optimization, or operational models.
6:00 GitHub Repo: A GitHub repository (MLP-from-scratch) is provided for accessing the code, although live commits were hindered by technical issues.
7:22 Business Use Case: The driving business problem is providing personalized predictions for different users across various use cases.
8:44 Design Principles: The tutorial uses test-driven development (TDD) combined with domain-driven design (DDD), highlighting ubiquitous language, bounded contexts, and aggregates.
9:47 Ubiquitous Language: Code should use language that aligns with business processes for clear communication with domain experts.
10:20 Bounded Contexts: Different parts of the system have clear boundaries; terms might have different meanings in different contexts (e.g., "model" in inference vs. training).
11:43 Aggregates: Grouping related objects (e.g., use case, user ID, model, feature vector) is central to structuring the domain model.
12:53 Starting with Inference: The build process begins with the inference engine, taking a TDD approach by writing a test for the simplest prediction (a constant value).
19:06 Defining Shared Types:shared_types.py is created to define common domain objects like User ID and UseCase for clarity and consistency.
21:21 Testing the Simplest Case: A test is written to assert that a constant prediction is returned, establishing the basic interface.
25:33 Supporting Different Use Cases: A second test is added to challenge the system to return different predictions for different use cases, pushing for a more complex implementation.
27:10 Evolution to Class-Based Inference: The inference engine is converted into a class to hold state and dependencies.
28:20 Introducing Models: A models.py file is created to define abstract interfaces for models (initially ConstantModel and later RulesBasedModel), enforcing a predict method.
30:00 Model Repository: The concept of a model repository is introduced as the way the inference engine accesses models, decoupling it from direct storage implementation.
32:18 Repository Pattern: The repository pattern is explained as an abstraction layer for data access (like models or features), enabling domain logic to be independent of external infrastructure details (e.g., GCS, Redis).
35:10 Defining Repository Interface: An abstract ModelRepository class is defined with a get_model method.
37:34 In-Memory Repository Implementation: An InMemoryModelRepository is created as a simple, testable concrete implementation using a dictionary. Custom exceptions (ModelNotFound) are introduced for clearer error handling.
40:34 Custom Exceptions:exceptions.py is added to define domain-specific exceptions.
41:39 Refactoring Inference Engine: The inference engine is updated to use the ModelRepository interface, removing the direct dictionary dependency.
43:00 Updating Tests: Existing tests are modified to work with the class-based inference engine and the injected in-memory model repository.
51:11 Predictions for Different Users: The next challenge is to make predictions that vary based on the user, requiring the introduction of features.
52:54 Introducing Features and Feature Vectors:shared_types.py is updated to define Feature, FeatureVector, FeatureSet, and FeatureSets to structure user-specific feature data.
56:03 Rules-Based Model: A RulesBasedModel is introduced which takes a FeatureVector as input to make a prediction, building upon the abstract model interface.
57:48 Feature Repository: Similar to the model repository, a FeatureRepository is introduced as an abstraction for accessing user features, with an InMemoryFeatureRepository implementation.
1:01:45 Integrating Feature Repository: The inference engine is updated to depend on a FeatureRepository to fetch features before making a prediction, now using both use case and user ID.
1:03:16 Testing with Features: A test is added to demonstrate that the inference engine can use features from the injected InMemoryFeatureRepository and a RulesBasedModel to make different predictions for different users.
1:05:49 Value of Architecture: The value of this decoupled, injectable architecture is highlighted for making the complex ML platform system highly testable and independent of specific infrastructure choices.
1:09:00 System Level View (Event-Driven Design): A larger code example is presented to show how multiple decoupled services (inference engine, feature store, model registry) can communicate.
1:11:48 Decoupled Communication: Since services are decoupled, they need a way to communicate, leading to event-driven design.
1:12:42 Event-Driven Design: Services don't communicate directly but through events and a central orchestrator (message bus) that consumes events and decides the next actions. Services raise domain exceptions when they can't complete a task.
1:15:00 Message Types: Three types of messages are defined: Command (what to do), Event (what happened), and Error (what went wrong).
1:16:30 Handlers and Translators: Handlers interpret domain exceptions into messages. Translators take events or errors and translate them into the next command to be executed.
1:19:37 Message Bus Orchestration: The message bus manages a queue of commands, dispatches them to handlers, translates responses into subsequent commands, and handles message prioritization (errors resolved first).
1:21:19 End-to-End Testing: The architecture enables testing the entire system flow (e.g., requesting prediction -> model/features not found errors -> message bus triggers fetching from repositories -> repositories publish events -> message bus triggers adding to inference engine's internal repos -> retries prediction successfully).
1:25:41 Challenge Questions: The tutorial concludes with challenges for further exploration, including integrating model training, hyperparameter tuning, model comparison, handling different hardware/runtimes, and using real message queues (Kafka, Pub/Sub) while maintaining testability.
1:27:47 Key Takeaway: The most important aspect of building an ML platform is thinking about how to test it, and DDD/EDD principles provide tools for creating testable, decoupled systems.
1:28:15 Conclusion: Nathan thanks the audience and encourages interaction via GitHub or LinkedIn for further discussion on these concepts.