NoOps Machine Learning. A PaaS End-to-End ML Setup with... | by Jacopo Tagliabue | Towards Data Science

“The princess you are looking for is in another castle.”
The rapid adoption of Machine Learning, from Big Tech to literally everyone else, resulted in a blooming season for ML tooling and what cool kids now call “MLOps” (see the thoughtful overview by Chip Huyen, if you need to catch up). While my LinkedIn feed obsessively repeats that “80% of machine learning projects don’t make it to production”, I feel that there actually dozens of strategies to deploy your models, but you’re kind of in the dark after you solved that: I am not talking about monitoring, but the general experience of how you develop, experiment, train, iterate in an effective way (i.e. if you still need to solve the production stuff, this post won’t really help).
Let’s face it, most ML scripts (well, certainly mine) are pretty messy, especially in the prototype-and-test-live-at-small-scale phase, as they involve stitching together several small tasks: retrieving data, training neural networks, running functional and behavioral tests etc. In other words, ML projects are DAG-like, and they need to be replayable, versioned, scheduled and so on: while we have successfully used Luigi before to orchestrate tasks, we hit some limitations with a growing team and big(ish) neural models.
There is an important phase in ML-driven product development where we should be able to trust our code, but retaining the possibility of making corrections as we go, without worrying about the overall production environment:
before full-scale engineering, ML needs a phase of good-enough engineering, where developers are somehow empowered end-to-end to ship their work to customers, internal users, etc.
This post is our <10 minutes introduction to a better developer life: an end-to-end ML project (from raw data to a working API) built around a simple philosophy — effective iterations means abstracting away all infrastructure from developers. While custom-built solutions may work well for special cases, we show how off-the-shelf tools — Metaflow, SageMaker, Serverless — can be combined for a great PaaS experience, as all computing is managed automatically for us:
in other words, from MLOps to NoOps at all.
All the code is shared on Github: clone the repo and tag along!
DISCLAIMER: this is not a tutorial on Metaflow, but a practical example of how tools can be combined to provide a better experience: while the repo should work with just copy+paste, some things may be unclear if you are not too familiar with Metaflow, SageMaker, etc.. It should be pretty straightforward to fill the gaps, but going through Netflix tutorials is always a good idea.
We assume the following setup on your dev machine (see the Medium story replies, for some additional tips): there is already a ton of materials on these steps and we want to focus on the value provided by the system as a whole (which is what justifies going through setup in the first phase!).
An AWS account: this is obvious, but please note that the free-tier won’t cut it for Metaflow! Make sure to have a suitable IAM role ready, as your development machine will need permissions to spin-up SageMaker endpoints.
Metaflow up and running: go through the AWS-based setup, which is simplified by the CloudFormation template! Please note that if the @batch decorator is commented out, a local setup may work as well with some changes: if you are just curious, you can start local and go to the full AWS-backed configuration after you’re sold on the approach.
Serverless up and running: make sure you can deploy code to AWS lambda from your terminal.
If you wish to add experiment tracking, we have added some tips at the end: to start, just open a community account for MLFlow or a personal account on Weights & Biases.
While not a definitive list, there used to be some frequently-mentioned pain points in our team when it comes to our everyday life as ML engineers:
We want to organize the small tasks composing a project into a modular and replayable DAG-like structure: we want Python and data objects to be re-usable across tasks and clearly versioned;
We want to assign computational resources effortlessly (i.e. without explicit provisioning/maintenance) per task, as we don’t want to run simple business logic on expensive GPUs, but we also don’t want to wait days for training runs. Since Snowflake allows us to do calculations directly in the data warehouse (a long story for another post!), our modus operandi is eminently non-distributed, and designed to be so: datasets fit comfortably in a laptop, and local development is way more efficient than iterating on Spark. Our bottleneck is mostly training, as we sometimes train big(ish) models on GPUs: however, naively switching context between local and cloud machines is very inefficient and time consuming;
We want to work well together, that is, having visibility on all our DAGs, experiments and artifacts, but also making sure we don’t overwrite each other’s runs by mistake;
We want a clear “path-to-make-an-impact”: even in a prototyping phase, we want support for deployment and scheduling — the goal is to able to push pretty robust micro-services without involving other teams.
Metaflow (+ SageMaker) allows us to satisfy all these desiderata in a unified framework (plus additional perks, like the super-performant s3 client!).
The project is a one-file DAG, centered around a regression model in Keras. The code is “toy” in many ways, but it shows how to go from a dataset to an API through a repeatable, modular and scalable process: the code is not particularly terse, but we aim for explanatory power over economy.
The DAG is depicted in the picture below, and can be broken down in four major parts:
Data loading: this step is achieved through a static file — thanks to Metaflow versioning, the exact dataset that is fed to training is automatically versioned.Model training: training the chosen model by exploring the hyperparameter space and measuring model performance on hold-out data.Model choice: collect all results from model runs, and chose the model for deployment based on some logic.Deployment: deploy the model to a SageMaker endpoint. Since SageMaker endpoints are internal to AWS, we use AWS lambda to make the predictions available through public APIs (see below).
Data loading: this step is achieved through a static file — thanks to Metaflow versioning, the exact dataset that is fed to training is automatically versioned.
Model training: training the chosen model by exploring the hyperparameter space and measuring model performance on hold-out data.
Model choice: collect all results from model runs, and chose the model for deployment based on some logic.
Deployment: deploy the model to a SageMaker endpoint. Since SageMaker endpoints are internal to AWS, we use AWS lambda to make the predictions available through public APIs (see below).
The entire script is ~200 lines of code and it should be pretty self-explanatory (if not, reach out!). The project showcases many features of Metaflow — such as file and object versioning, parameters, automatic retry, etc. — , but we focus here on two: parallelization and remote computing.
Parallelization: the key lines are the following, in which we first define a range of parameters (here, using different learning rates for pedagogical purposes) and then we use the foreach keyword:
self.learning_rates = [0.1, 0.2]self.next(self.train_model, foreach='learning_rates')
What happens is that the train_model step will be executed in parallel, with Metaflow spinning out new processes, each training a model with a particular learning rate.
Remote computing: the key line is the following (you can actually comment it out the first time, to appreciate how your development experience is really the same, with or without cloud-computing):
@batch(gpu=1, memory=80000)
With this decorator, you tell Metaflow that the ensuing task has some computing requirements (here expressed as GPU and memory capacity, but “number of CPUs” is a popular option): Metaflow automatically creates a job in AWS Batch, then runs the step on a GPU-enabled machine. This shows how to integrate local dataset preparation with remote training, without change in code or explicit resource provisioning: the DAG stays the same, and it is Metaflow job to distribute the computation.
Metaflow gets us the comfort and speed of local debugging in data wrangling, and then a pool of unbounded computational resources when horsepower is needed.
If you are curious to see a more realistic flow, you can appreciate below the power of parallelization and remote computing: you retrieve product images from the warehouse, speed-up vectorization, get back to your laptop for the final touches (e.g. filtering), and then finally train your deep neural network.
“Talk is cheap, show me the code” — L. Torvalds
To run the project, cd into the flow directory and type the following command under your Python interpreter (make sure packages in requirements.txt are installed!):
python training.py run
After some quick checks, the DAG will load the data and then start parallelizing the training steps (if you are running with the batch decorator, AWS Batch console will track the two processes for you):
At the end of each training process, the terminal will print the cloud path of the model, as versioned by the built-in s3 client:
At the end of the DAG, a new SageMaker endpoint will be deployed: write down its name, as that is the only information we need to start serving predictions through a public API.
After a run is concluded, we have to make sure our new model can be accessed through a devoted micro-service. We deploy a PaaS endpoint, leveraging the Infrastructure-as-Code setup in serverless.yml: simply cd into the serverless folder, and run the following command:
serverless deploy --sagemaker my-endpoint
where my-endpoint is the name of the endpoint created by the DAG. After a minute or so, you should be greeted by a success message:
Write down the endpoint URL, paste it in the browser and try to supply one x value:
https://my-url.amazonaws.com/dev/predict?x=9.45242
Et voilà — the browser should return our JSON response, with a pretty decent execution time:
As far as systematically tracking our experiments, we could exploit Metaflow built-in versioning and its handy client API to retrieve what we stored in:
self.hist = history.history
and quickly plot loss / epochs in a notebook — but in a truly PaaS spirit of never write code unless you really need to, why don’t we re-use the neat UI of MLflow or W&B? Luckily, we can add tracking to our DAG with a few lines of code.
At the very beginning of the train_model step, we inject some small code to make sure that i) the machine (local or remote) running the code has MLflow/W&B installed and ii) credentials are available as environment variables.
os.system('pip install mlflow==1.13.1')os.environ["DATABRICKS_HOST"] = 'https://community.cloud.databricks.com' # don't change thisos.environ["DATABRICKS_USERNAME"] = 'my_user_name'os.environ["DATABRICKS_PASSWORD"] = 'my_password'
(The same basic idea applies to W&B, but the env variable you need to set is the API KEY). Now that the library is installed, we can just exploit autologging to start tracking:
import mlflow# don't change this for Community editionmlflow.set_tracking_uri("databricks")# your experiment name mlflow.set_experiment("/Users/my_user_name/exp_name") mlflow.tensorflow.autolog()
The next time the DAG is executed, the Databricks dashboard will record and display the key metrics for all our runs:
This is the same output, using W&B instead:
As promised, we went from data loading to a model-based API in less than 10 minutes and without any infrastructure work: if you have the same pain points we had, our script gives you a worked-out example of how Metaflow-centric development (plus, some carefully chosen PaaS services) radically simplifies real-life machine learning projects.
While the NoOps slogan is certainly a bit too good to be literally true, we do believe that good abstractions will make DevOps operations less and less visible in ML pipelines, allowing us to focus more and more on data and models.
Alas, what is the history of software engineering, if not a series of powerful abstractions?
Don’t forget to check the additional notes (in the story replies) for more technical details, setup tips and notes for future iterations. If you have questions or feedback, please connect and share your MLOps story (all Medium opinions are my own).
Thanks to Christine, Ethan, Luca and Patrick for precious comments to earlier version of this post + repo.