Presented at All Things Open 2023 Presented by Danny McCormick - Google Title: Deploying Models at Scale with Apache Beam Abstract: Apache Beam is an open source tool for building distributed scalable data pipelines. This talk will explore how Beam can be used to perform common machine learning tasks, with a heavy focus on running inference at scale. The talk will include a demo component showing how Beam can be used to deploy and update models efficiently on both CPUs and GPUs for inference workloads. An attendee can expect to leave this talk with a high level understanding of Beam, the challenges of deploying models at scale, and the ability to use Beam to easily parallelize their inference workloads. Find more info about All Things Open: On the web: https://www.allthingsopen.org/ Twitter: https://twitter.com/AllThingsOpen LinkedIn: https://www.linkedin.com/company/all-things-open/ Instagram: https://www.instagram.com/allthingsopen/ Facebook: https://www.facebook.com/AllThingsOpen Mastodon: https://mastodon.social/@allthingsopen Threads: https://www.threads.net/@allthingsopen 2023 conference: https://2023.allthingsopen.org/

1. Deploying Models at Scale with
Apache Beam
Danny McCormick
docs.google.com/presentation/d/1WBcT69Lh9GIFoDhgx5OwuCJp0LkoCLxRPiU3Wa_67ns
(or shorturl.at/hAKN2)

2. Who am I

3. Me!

4. Topics
1. Beam History + Overview
2. Basic ML Inference
3. Problems/Solutions
○ Model Freshness
○ Large Models
○ Specialty Hardware
4. Where Next?

5. What is Apache Beam

6. Data got big:

7. Really big:
Monday
Tuesday
Wednesday
Thursday
Friday

8. Bigger! And neverending!
9:00
8:00 14:00
13:00
12:00
11:00
10:00

9. In the beginning, there was MapReduce
Datastore
Map
Map
Map
Map
Map
Map
Reduce
Reduce
Reduce
Reduce
Reduce
Reduce
Shuffle
Datastore

10. In the beginning, there was MapReduce

11. Then came Flume
Datastore
Map
Map
Datastore
Map
Group by Key
(Reduce)
Combine
Map
Map
Combine
Datastore
Datastore
Datastore

12. From Flume came Beam
Datastore
Map
Map
Datastore
Map
Group by Key
(Reduce)
Combine
Map
Map
Combine
Datastore
Datastore
Datastore

13. Unified Model for Batch and Streaming
● Batch processing is a special case of
stream processing
● Batch + Stream = Beam
Mond
ay
Tues
day
Wedn
esdayThurs
day
Friday

14. Build your pipeline in whatever language(s) you want…
Group by Key

15. … with whatever execution engine you want
Cloud Dataflow
Apache Spark
Apache Flink
Dask
Ray
Apache Samza
Prism/Local
Twister2
Group by Key

16. Beam Basics

17. Terms
● PCollection - distributed multi-element
dataset
● Transform - operation that takes N
PCollections and produces M PCollections
● Pipeline - directed acyclic graph of
Transforms and PCollections

18. Basic Beam Graph
Source
Transform
Sink
Transform
Source
Transform
Map
Transform
Combine
Transform
Sink
Transform
Sink
Transform

19. Basic Beam Pipeline
def add_one(element):
return element + 1
import apache_beam as beam
with beam.Pipeline() as pipeline:
pipeline
| beam.io.ReadFromText('gs://some/inputData.txt')
| beam.Map(add_one)
| beam.io.WriteToText('gs://some/outputData')
Read Text file
Map
Transform
Write to text
file
0
15
3
21
1
16
4
22

20. Beam ML

21. The ML lifecycle

22. Inference with Beam

23. Challenges of Distributed Inference
● Efficiently loading models
● Batching
● Using multiple models

24. Distributed Inference with Beam
● Beam takes care of all of this with the
RunInference transform
● Loads model, batches inputs, and
plugs into DAG
RunInference(model_handler=<config>)

25. RunInference
>>> model_handler = PytorchModelHandlerTensor(
... model_class=LinearRegression,
... model_params={'input_dim': 1, 'output_dim': 1},
... state_dict_path='gs://path/to/model.pt')
>>> with beam.Pipeline() as p:
... predictions = (
... p
... | beam.Create(numpy.array([10, 40, 60, 90],
... dtype=numpy.float32).reshape(-1, 1))
... | beam.Map(torch.Tensor) # Map np array to Tensor
... | RunInference(model_handler=model_handler)
... | beam.Map(print))

26. Basic Inference Demo
colab.sandbox.google.com/github/apache/beam/blob/master/examples/notebooks/beam-ml/run_inference_huggingface.ipynb
(shorturl.at/brvN9)

27. Automatic Model Refresh

28. You’ve deployed your model! Now what?
● New data
● New training algorithms
● New models

29. How can we update without bringing down Prod?

30. Automatic Model Refresh
● Hot swaps model in live pipeline
● Manages memory for you
● No pipeline down time (though
maybe some inference down time)

31. Automatic Model Refresh
side_input_pcoll = (pipeline
| "WatchFilePattern" >> WatchFilePattern(file_pattern=file_pattern))
inferences = (image_data
| "RunInference" >> RunInference(model_handler=model_handler,
model_metadata_pcoll=side_input_pcoll))

32. Large Models

33. Distributed Runner Architecture*
VM
Worker
Process
Worker
Process
Worker
Process
Worker
Process
Worker
Process
Worker
Process
Worker
Process
Worker
Process

34. Ideal small model configuration
VM
I/O &
Model
I/O &
Model
I/O &
Model
I/O &
Model
I/O &
Model
I/O &
Model
I/O &
Model
I/O &
Model

35. Ideal Large Model Configuration
VM
I/O
I/O
I/O
I/O
I/O
I/O
I/O
I/O
Model

36. Ideal Multi Large Model Configuration
VM
I/O
I/O
I/O
I/O
I/O
I/O
I/O
I/O
Model 1
VM
I/O
I/O
I/O
I/O
Model 2
I/O
I/O
I/O
I/O
Model 3
Option 1 Option 2

37. How do we map ideal model configurations to this?
VM
Worker
Process
Worker
Process
Worker
Process
Worker
Process
Worker
Process
Worker
Process
Worker
Process
Worker
Process

38. Ideal small model configuration
VM
I/O &
Model
I/O &
Model
I/O &
Model
I/O &
Model
I/O &
Model
I/O &
Model
I/O &
Model
I/O &
Model

39. Default configuration: share model across threads
>>> model_handler = PytorchModelHandlerTensor(
... model_class=LinearRegression,
... model_params={'input_dim': 1, 'output_dim': 1},
... state_dict_path='gs://path/to/model.pt')
>>> pcoll | RunInference(model_handler=model_handler)

40. Ideal Large Model Configuration
VM
I/O
I/O
I/O
I/O
I/O
I/O
I/O
I/O
Model

41. Ideal Large Model Configuration
VM
Worker
Process
Worker
Process
Worker
Process
Worker
Process
Worker
Process
Worker
Process
Worker
Process
Worker
Process
Inference
Process

42. Optional: serve a single model for all processes
>>> model_handler = PytorchModelHandlerTensor(
... model_class=LinearRegression,
... large_model=True,
... model_params={'input_dim': 1, 'output_dim': 1},
... state_dict_path='gs://path/to/model.pt')
>>> pcoll | RunInference(model_handler=model_handler)

43. Ideal Multi Large Model Configuration
VM
I/O
I/O
I/O
I/O
I/O
I/O
I/O
I/O
Model 1
VM
I/O
I/O
I/O
I/O
Model 2
I/O
I/O
I/O
I/O
Model 3

44. Ideal Large Model Configuration
VM
Worker
Process
Worker
Process
Worker
Process
Worker
Process
Worker
Process
Worker
Process
Worker
Process
Worker
Process
Model
Manager

45. Optional: serve a single model for all processes
>>> per_key_mhs = [
... KeyModelMapping(['key1', 'key2', 'key3'], model_handler_1),
... KeyModelMapping(['foo', 'bar', 'baz'], model_handler_2)]
>>> mh = KeyedModelHandler(per_key_mhs)
>>> pcoll | RunInference(model_handler=mh)

46. Large Model Demo
colab.sandbox.google.com/github/apache/beam/blob/master/examples/notebooks/beam-ml/per_key_models.ipynb
(shorturl.at/pKNU2)

47. Specialty Hardware

48. GPU/TPU Support
● Hardware availability dependent on
runner
● Beam has some primitives that help

49. Beam Primitives for GPUs
● Resource hints for heterogeneous
pools
● Built in detection + framework
specific responses to GPUs at the
ModelHandler level
● Large model setting (revisited)

50. Ideal Large Model Configuration
VM
Worker
Process
Worker
Process
Worker
Process
Worker
Process
Worker
Process
Worker
Process
Worker
Process
Worker
Process
Inference
Process

51. Where next?
(opportunities I see, not representative of the whole community)

52. Inference Space
● More frameworks
● Better performance testing/profiling
● Model Manager Improvements

53. Beyond Inference
● MLTransform for data prep and
pre/postprocessing
● Feature Store Enrichment
● Higher level ML support (e.g.
anomaly detection)

54. Come join our community!

55. Questions?
Contact - Danny McCormick (dannymccormick@google.com)
Slides - shorturl.at/hAKN2