Loading…
Training Orchestrator: Unifying Model Training at Yelp
YelpYing Wang and Nathan Sponberg, Software Engineer
Summary
Yelp developed Training Orchestrator to standardize machine learning pipeline execution and eliminate divergent training scripts across its applied engineering teams. Previously, monolithic Spark scripts coupled training logic directly to cluster runtimes, preventing local testing, slowing debug cycles, and causing inconsistent configuration drift across projects. The new framework introduces a declarative configuration layer backed by Pydantic models, enabling static type checking and step schema validation before runtime execution starts. Training workflows are defined as modular step hierarchies that the orchestration engine compiles into directed acyclic graphs for execution in topological order. By injecting Spark and MLflow contexts into steps, the architecture ensures reproducible tracking, automated Slack alerts, and parity across local development and production environments.
Context
Applied machine learning teams at Yelp maintained separate Spark-based training batches, scripts, and configurations. Over time, these diverged, resulting in code duplication, configuration drift, scattered validation checks, and fragile custom scripts for monitoring and notifications. Monolithic training scripts were tightly coupled to Spark cluster runtimes and job submission mechanisms, preventing local runs and making unit and integration testing difficult.
Approach / What changed
Yelp built Training Orchestrator, a declarative, configuration-driven platform that separates training code from execution infrastructure. Pipelines are defined using Pydantic configurations that enforce type safety and validate step input and output schemas upfront. The orchestrator automatically builds a directed acyclic graph from declared step dependencies, executes steps in topological order, injects Spark and MLflow execution contexts, and standardizes Slack notifications.
Takeaways
- Decoupling training logic from Spark cluster submission allows developers to run identical pipeline configurations locally or in Jupyter, making unit testing and local debugging feasible.
- Training Orchestrator uses Pydantic to validate orchestrator parameters and step input/output schemas at instantiation, catching configuration mismatches before initiating distributed compute jobs.
- Defining workflows declaratively lets the engine construct a directed acyclic graph that enforces prerequisite completion, topological step execution, and shared MLflow tracking across all steps.
Related reading
Yelp ·
Beyond the Menu Tree: How Yelp Built a Smarter Customer Success Chatbot with AI
Yelp transitioned from a legacy customer support chatbot that relied on rigid phrase matching and static menus to a retrieval-augmented generation system. The updated architecture classifies incoming queries using a large language model to route users into one of five specialized workflows, including QA, Billing, Refund, Cancel, and Review. For general QA, the system uses FAISS to search an in-memory vectorstore built exclusively from article metadata such as titles, summaries, and headers rather than raw article text. Retrieved articles are supplied to the language model as prompt context, followed by automated validations for safety, character limits, and hallucinated hyperlinks. In production A/B testing, this metadata-driven design doubled the chatbot resolution rate while achieving approximately 94% recall@5.
Lina Lee, Machine Learning Engineer; Nelson Lee, Engineering ManagerYelp ·