This document summarizes a talk on building an ML platform with Ray and MLflow. Ray is an open-source framework for distributed computing and machine learning. It provides libraries like Ray Tune for hyperparameter tuning and Ray Serve for model serving. MLflow is a tool for managing the machine learning lifecycle including tracking experiments, managing models, and deploying models. The talk demonstrates how to build an end-to-end ML platform by integrating Ray and MLflow for distributed training, hyperparameter tuning, model tracking, and low-latency serving.

1. Building an ML
Platform with Ray and
MLflow
Amog Kamsetty and Archit Kulkarni
Ray Team @ Anyscale

2. The Team
Archit Kulkarni Amog Kamsetty Dmitri Gekhtman Edward Oakes
Richard Liaw Kai Fricke Simon Mo
Kathryn Zhou

3. Overview of Talk
▪ What are ML Platforms?
▪ Ray and its libraries
▪ MLflow
▪ Demo: An ML Platform
built with MLflow and
Ray

4. What are ML Platforms?

5. Typical ML Process
Fuzzy
search!
NLP, DL …

6. Execution
- Feature engineering
- Training
- Including tuning
- Serving
- Offline scoring, inference
- Online serving
Typical ML Process -- Simplified
Management
- Tracking
- Data, Code, Configurations
- Reproducing Results
- Deployment
- Deploy in a variety of
environments

7. Challenges with the ML Process
Data/Features
• Data Preparation
• Data Analysis
• Feature
Engineering
• Data Pipeline
• Data
Management/Feat
ure Store
• Manages big data
clusters
Model
• ML Expertise
• Implement SOTA
ML Research
• Experimentation
• Manage GPU
infrastructure
• Scalable training &
hyperparameter
tuning
Production
• A/B Testing
• Model Evaluation
• Analysis of
Predictions
• Deploy in variety of
environments
• CI/CD
• Highly Available
prediction service
Data/Research
Scientist
Engineers

8. Challenges with the ML Process
Data
• Data Preparation
• Data Analysis
• Feature
Engineering
• Data Pipeline
• Data
Management/Feat
ure Store
• Manages big data
clusters
Model
• ML Expertise
• Implement SOTA
ML Research
• Experimentation
• Manage GPU
infrastructure
• Scalable training &
hyperparameter
tuning
Production
• A/B Testing
• Model Evaluation
• Analysis of
Predictions
• Deploy in variety of
environments
• CI/CD
• Highly Available
prediction service
Data/Research
Scientist
Software/Data/
ML Engineer
ML Platform
Abstraction

9. ML Platforms -- Scale
- LinkedIn:
- 500+ “AI engineers” building models; 50+ MLP engineers
- > 50% offline compute demand (12K servers each with 256G RAM)
- More than 2x a year
- Uber Michelangelo, AirBnB Bighead, Facebook FBLearner,
etc.
- Globally, a few Billion $ now, growing 40%+ YoY
- Many companies building ML Platforms from the ground up

10. ML Platforms -- Landscape
(Source: Intel Capital)

11. ML Platforms -- Landscape
(Source: Intel Capital)

12. Execution
- Feature engineering 🔪
- Training 🍳
- Including tuning 🧂
- Serving 🍽
- Offline scoring, inference
- Online serving
Typical ML Process -- Simplified
Management
- Tracking 📝
- Data, Code, Configurations
- Reproducing Results 📖
- Deployment 🚚 💻
- Deploy in a variety of
environments

13. Execution
- Feature engineering 🔪
- Training 🍳
- Including tuning 🧂
- Serving 🍽
- Offline scoring, inference
- Online serving
Typical ML Process -- Simplified
Management
- Tracking 📝
- Data, Code, Configurations
- Reproducing Results 📖
- Deployment 🚚 💻
- Variety of environments

14. Ray and its Libraries

15. What is Ray?
• A simple/general library for distributed computing
• Single machine or 100s of nodes
• Agnostic to the type of work
• An ecosystem of libraries (for scaling ML and more)
• Native: Ray RLlib, Ray Tune, Ray Serve
• Third party: Modin, Dask, Horovod, XGBoost, Pytorch Lightning
• Tools for launching clusters on any cloud provider

16. Three key ideas
Execute remote functions as tasks, and
instantiate remote classes as actors
• Support both stateful and stateless computations
Asynchronous execution using futures
• Enable parallelism
Distributed (immutable) object store
• Efficient communication (send arguments by reference)

17. Ray API

18. API
Functions -> Tasks
def read_array(file):
# read array “a” from “file”
return a
def add(a, b):
return np.add(a, b)

19. API
Functions -> Tasks
@ray.remote
def read_array(file):
# read array “a” from “file”
return a
@ray.remote
def add(a, b):
return np.add(a, b)

20. API
Functions -> Tasks
@ray.remote
def read_array(file):
# read array “a” from “file”
return a
@ray.remote
def add(a, b):
return np.add(a, b)
id1 = read_array.remote(“/input1”)
id1
read_array

21. API
Functions -> Tasks
@ray.remote
def read_array(file):
# read array “a” from “file”
return a
@ray.remote
def add(a, b):
return np.add(a, b)
id1 = read_array.remote(“/input1”)
id2 = read_array.remote(“/input2”)
id1
read_array
id2
zeros
read_array

22. API
Functions -> Tasks
@ray.remote
def read_array(file):
# read array “a” from “file”
return a
@ray.remote
def add(a, b):
return np.add(a, b)
id1 = read_array.remote(“/input1”)
id2 = read_array.remote(“/input2”)
id3 = add.remote(id1, id2)
id1
read_array
id2
zeros
read_array
id3
add

23. API
Functions -> Tasks
@ray.remote
def read_array(file):
# read array “a” from “file”
return a
@ray.remote
def add(a, b):
return np.add(a, b)
id1 = read_array.remote(“/input1”)
id2 = read_array.remote(“/input2”)
id3 = add.remote(id1, id2); ray.get(id3)
id1
read_array
id2
zeros
read_array
id3
add

24. API
Functions -> Tasks
@ray.remote
def read_array(file):
# read array “a” from “file”
return a
@ray.remote
def add(a, b):
return np.add(a, b)
id1 = read_array.remote(“/input1”)
id2 = read_array.remote(“/input2”)
id3 = add.remote(id1, id2)
Classes -> Actors

25. API
Functions -> Tasks
@ray.remote
def read_array(file):
# read array “a” from “file”
return a
@ray.remote
def add(a, b):
return np.add(a, b)
id1 = read_array.remote(“/input1”)
id2 = read_array.remote(“/input2”)
id3 = add.remote(id1, id2)
Classes -> Actors
@ray.remote
class Counter(object):
def __init__(self):
self.value = 0
def inc(self):
self.value += 1
return self.value

26. API
Functions -> Tasks
@ray.remote
def read_array(file):
# read array “a” from “file”
return a
@ray.remote
def add(a, b):
return np.add(a, b)
id1 = read_array.remote(“/input1”)
id2 = read_array.remote(“/input2”)
id3 = add.remote(id1, id2)
Classes -> Actors
@ray.remote
class Counter(object):
def __init__(self):
self.value = 0
def inc(self):
self.value += 1
return self.value
c = Counter.remote()
id4 = c.inc.remote()
id5 = c.inc.remote()
ray.get([id4, id5])

27. API
Functions -> Tasks
@ray.remote
def read_array(file):
# read array “a” from “file”
return a
@ray.remote(num_gpus=1)
def add(a, b):
return np.add(a, b)
id1 = read_array.remote(“/input1”)
id2 = read_array.remote(“/input2”)
id3 = add.remote(id1, id2)
Classes -> Actors
@ray.remote(num_gpus=1)
class Counter(object):
def __init__(self):
self.value = 0
def inc(self):
self.value += 1
return self.value
c = Counter.remote()
id4 = c.inc.remote()
id5 = c.inc.remote()
ray.get([id4, id5])

28. at Anyscale
Your app
here!
Native Libraries 3rd Party Libraries
Ecosystem
Universal framework for
Distributed computing
Ray Ecosystem

29. Ray Tune

30. Ray Tune: Scalable
Hyperparameter Tuning
Wide variety of algorithms Compatible with ML frameworks
HYPERBAND
PBT
BAYESIAN OPT.

31. Ray Tune focuses on
simplifying execution
Easily launch distributed multi-gpu
tuning jobs
Automatic fault tolerance to save
3x on GPU costs
https://www.vecteezy.com/
$ ray up {cluster config}
ray.init(address="auto")
tune.run(func, num_samples=100)

32. Ray Tune interoperates
with other HPO libraries
Ray Tune
Ax
Optuna
scikit-optimize
…

33. def train_model(config={}):
model = ConvNet(config)
for i in range(steps):
current_loss = model.train()

34. from ray import tune
def train_model(config={}):
model = ConvNet(config)
for i in range(steps):
current_loss = model.train()
tune.report(loss=current_loss)

35. def train_model(config):
model = ConvNet(config)
for i in range(epochs):
current_loss = model.train()
tune.report(loss=current_loss)
tune.run(train_model,
config={“lr”: 0.1})

36. tune.run(
train_model,
config={“lr”: tune.uniform(0.001, 0.1)},
num_samples=100
)
def train_model(config):
model = ConvNet(config)
for i in range(epochs):
current_loss = model.train()
tune.report(loss=current_loss)

37. tune.run(
train_model,
config={“lr”: tune.uniform(0.001, 0.1)},
num_samples=100,
scheduler=ASHAScheduler())
def train_model(config):
model = ConvNet(config)
for i in range(epochs):
current_loss = model.train()
tune.report(loss=current_loss)

38. tune.run(
train_model,
config={“lr”: tune.uniform(0.001, 0.1)},
num_samples=100,
scheduler=PopulationBasedTraining(...))
def train_model(config, checkpoint_dir=None):
model = ConvNet(config)
if checkpoint_dir is not None:
model.load_checkpoint(checkpoint_dir+”model.pt”)
for i in range(epochs):
current_loss = model.train()
with tune.checkpoint_dir() as dir:
model.save_checkpoint(dir+”model.pt”)
tune.report(loss=current_loss)

39. Ray Serve

40. Ray Serve is a
Web Framework
Built for
Model Serving

41. Model Serving in Python

42. Ray Serve is
high-performance and flexible
• Framework-agnostic
• Easily scales
• Supports batching
• Query your endpoints from
HTTP and from Python
• Easily integrate with other
tools

43. Ray Serve is built on top of Ray
For user, no need to think about:
• Interprocess communication
• Failure management
• Scheduling
Just tell Ray Serve to scale up your model.

44. Serve functions and stateful classes.
Ray Serve will use multiple replicas to parallelize
across cores and across nodes in your cluster.
Ray Serve API

45. Flexibility
Query your model from HTTP:
> curl "http://127.0.0.1:8000/my/route"
Or query from Python using ServeHandle:

46. MLflow

47. Challenges of ML in production
• It’s difficult to keep track of experiments.
• It’s difficult to reproduce code.
• There’s no standard way to package and deploy
models.
• There’s no central store to manage models (their
versions and stage transitions).
Source: mlflow.org

48. What is MLflow?
• Open-source ML lifecycle management tool
• Single solution for all of the above challenges
• Library-agnostic and language-agnostic
• (Works with your existing code)

49. Four key functions of MLflow
Source: MLflow

50. MLflow Tracking

51. MLflow Models

52. Ray + MLflow

53. Ray Tune + MLflow Tracking
def train_model(config):
model = ConvNet(config)
for i in range(epochs):
current_loss = model.train()
tune.report(loss=current_loss)
tune.run(
train_model,
config={“lr”: tune.uniform(0.001, 0.1)},
num_samples=100,
callbacks=[MLflowLoggerCallback(“my_experiment”)])

54. Ray Tune + MLflow Tracking
@mlflow_mixin
def train_model(config):
mlflow.autolog()
xgboost_results = xgb.train(config, ...)
tune.run(
train_model,
config={“lr”: tune.uniform(0.001, 0.1)},
num_samples=100)

55. +
> pip install mlflow-ray-serve
> ray start --head
> serve start

56. MLflow deployments CLI
Create deployment
> mlflow deployments create -t ray-serve -m <model URI>
--name my_model -C num_replicas=100
Model URI:
• models:/MyModel/1
• runs:/93203689db9c4b50afb6869
• s3://<bucket>/<path>
• ...

57. MLflow deployments Python API
Create model

58. Integrating with Ray Serve is easy.
• Ray Serve endpoints can be called from Python.
• Clean conceptual separation:
• Ray Serve handles data plane (processing)
• MLflow handles control plane (metadata, configuration)

59. Demo: An ML Platform built with MLflow and Ray

60. Acknowledgements
Thanks to Jules Damji, Sid Murching, and Paul Ogilvie for
their help and guidance with MLflow.
Thanks to Dmitri Gekhtman, Kai Fricke, Simon Mo,
Edward Oakes, Richard Liaw, Kathryn Zhou and the rest
of the Ray team!

62. Feedback
Your feedback is important to us.
Don’t forget to rate and review the sessions.