Yelp Launches Training Orchestrator for Unified Machine Learning Model Training

Introduction of Training Orchestrator

Yelp has launched Training Orchestrator, an internal framework aimed at unifying machine learning model training. This new system replaces individual team Spark training scripts, addressing inefficiencies and inconsistencies that previously hampered the training process.

Challenges with Previous System

Prior to Training Orchestrator, each machine learning team at Yelp had developed its own orchestration scripts. This led to issues such as duplicated code, inconsistent configurations, and difficulties with monitoring. The previous framework involved monolithic scripts tied to Spark clusters, which made local testing challenging and slowed down the iteration process.

Benefits of the New Framework

The Training Orchestrator separates the 'what' from the 'how' in the execution of training pipelines. Pipelines can now utilize Pydantic-based configuration objects, making them more manageable and error-resistant.

Developers can validate their configurations upfront, significantly reducing wasted compute resources. This approach also allows teams to run models simultaneously and reuse functions without altering the orchestration code.

Enhanced Testing and Validation

With the new framework, schema validation catches mismatches quickly, preventing potential hours of wasted compute during a running job. Developers are able to run full pipelines locally and write unit tests for each step, enhancing their overall productivity. By caching inputs for data-loading steps, integration testing is further expedited.

Logging and Tracking Improvements

The Training Orchestrator automatically logs the complete configuration of every run as an MLflow artifact. This feature facilitates reproducibility, enabling users to rerun previous configurations from a single recorded source. Additionally, the framework supports real-time notifications and tracking, which enhances visibility into the training processes.