From Machine Learning Code to MLOps – A Reference Architecture

Amazon SageMaker is a fully managed service that makes it easy for enterprises to build end-to-end production-ready machine learning pipelines without sacrificing speed, security, and accessibility. This article will propose a reference architecture based on the Amazon SageMaker ecosystem so you can get started right away with your own ML projects operating on the AWS platform. 


Fig 1: MLOps components (image by the author)

We will start with the simplest form of Machine Learning Operations (MLOps) and gradually add other building blocks to have a complete picture in the end. Let’s dive in!

Exploration block:

Given a business problem or a process improvement opportunity identified and documented by the business analyst, the machine learning operation starts with exploratory data analysis “EDA” where data scientists familiarize themselves with a sample of data and apply several machine learning techniques and algorithms to find the best ML solution. They will leverage Amazon SageMaker Studio which is a web-based integrated development environment (IDE) for machine learning to ingest the data, perform data analysis, process the data, and train and deploy models for making inferences using a non-production endpoint. Inside Amazon SageMaker Studio they have access to Amazon SageMaker Data Wrangler which contains over 300 built-in data transformations to quickly prepare the data without having to write any code. Amazon Athena and AWS Glue are other tools that can be used to explore and prepare data. All the experiments by the data scientists will be tracked using SageMaker Experiment capability for reproducibility.

Fig 2: MLOps Exploration block (Image by the author)

ML Workflow block:

Next, machine learning engineers convert the proposed solution by the data scientist to the production-ready ML code and create end-to-end machine learning workflow including data processing, feature engineering, training, model evaluation, and model creation for deployment using a variety of available hosting options. In AWS there are 4 options for orchestrating end-to-end ML workflow with Amazon SageMaker integration: 

  • Amazon SageMaker Pipeline
    • Using Pipelines SDK a series of interconnected steps will build the entire ML pipeline that is defined using a directed acyclic graph (DAG).
  • Amazon Managed Workflow for Apache Airflow (MWAA)
    • Using Airflow SageMaker operators or using Airflow PythonOperator end-to-end ML pipeline can be configured, scheduled, and monitored.
  • AWS Step Function workflow
    • Integration between AWS SageMaker and AWS Step Functions through the AWS Step Function Data Science SDK allows one to easily create multi-step machine learning workflows.
  • Kubeflow Orchestration
    • SageMaker components for Kubeflow pipelines allows you to submit SageMaker processing, training, HPO jobs, and deploy the model directly from the Kubeflow pipeline workflow.

As part of the model evaluation/test step, AWS Step Function is launched to run a comprehensive suite of ML-related tests. Additionally, Amazon SageMaker Feature Store is used to store, share, and manage features for machine learning (ML) models during training (offline storage) and inference (online storage). Finally, SageMaker ML Lineage tracking is enabled to track data, and model lineage metadata which is crucial for ML workflow reproducibility, model governance, and audit standards. 

Fig. 3: MLOps ML Workflow block (image by the author)

Continuous Integration block:

After implementing the first two blocks we already have a fully functioning ML workflow and endpoint that can be called by users or applications to consume the ML prediction. To take this to the next level and to eliminate manual work as a result of making any update to the code or infrastructure, an MLOps engineer will build up the Continuous Integration (CI) block enabling data scientists or ML engineers to regularly merge their changes into AWS CodeCommit, after which automated builds and tests are run using AWS CodeBuild including building any custom ML image based on the latest container hosted on Amazon ECR which will be referenced by the ML workflow consequently.

Fig. 4: MLOps Continuous Integration block (Image  by the author)

Continuous Deployment block:

Not only we want to build and test our ML application each time a code change is pushed to AWS CodeCommit, but also, we want to deploy the ML application in production continuously. In the Continuous Deployment (CD) block the proposed solution is extended to decouple the model training workflow from the model deployment section. This provides an opportunity to update the model configuration and infrastructure without affecting the training workflow. To manage model versions, model metadata, and automate the model deployment with CI/CD, the Amazon SageMaker Model Registry is used. Here, the manual approver (e.g. the senior data scientist), will approve the model which triggers a new automated deployment through CodePipeline, CodeBuild, and AWS CloudFormation.. The Automatic Rollback feature of AWS CloudFormation is key here, in case anything goes wrong during the deployment. SageMaker elastic endpoint autoscaling is also configured to handle changes in the inference workload. 

Fig. 5: MLOps Continuous Deployment block (Image by the author)

Monitoring & Continuous Training block:

When building machine learning models we assume that the data that the model was trained on will be similar to the data used when making predictions. Therefore, any drift in data distribution or model performance on the new inference data will result in data scientists needing to revisit the features or retrain the model to reflect the most recent changes. The Continuous Training (CT) block continuously monitors  the data and model quality to detect any bias or drift and inform a timely remediation . This will be achieved by enabling Amazon SageMaker Model Monitor in the inference endpoint configuration, creating baseline jobs during data prep and training along with SageMaker Clarify feature and Amazon CloudWatch event to automatically monitor for any drift, trigger re-training, and notify the relevant parties.

Fig. 6: MLOps Continuous Training block (Image by the author)

Security and governance consideration: 

To make our MLOps architecture secure, some important security and governance requirements are recommended. This includes: 

  • Deploying the solution in a VPC and securely connecting resources to SageMaker Studio through VPC endpoints.
  • Separation of concerns through multi-account architecture. AWS Control Tower offers a mechanism to easily set up and secure multi-account AWS environments.
  • Defining proper user and service roles to perform required operations on the AWS account and running tasks such as SageMaker training and ML workflow. AWS IAM is a powerful service for this purpose.
  • Implementing the least-privilege policy. AWS Service Catalog allows customers to implement products with required resources centrally provisioned without requiring each user to launch them separately.
  • Enforcing guardrails such as requiring proper resource tagging or limiting the type of resources used, for different users and roles. AWS Organizations and its Service Control Policies (SCP) feature can be used as an enterprise scale guardrail management.
  • Encrypting the traffic at rest and in transit. Beyond the default encryption of model artifacts and storage attached to training instances by SageMaker, AWS KMS key can be used to encrypt sensitive data.

In this blog post, we have proposed a reference architecture to build an end-to-end secure, scalable, repeatable, and reproducible MLOps solution to process data, train, create & update models, as well as deploy and monitor deployed model using SageMaker features with CI/CD/CT concepts incorporated. 

If your team needs expert assistance to deploy Machine Learning models on AWS at scale, consider engaging with Caylent to craft a custom roadmap with our collaborative MLOps Strategy Assessment and/or realize your vision through our MLOps pods.

Share this article

Leave a comment


Share this article


Join Thousands of DevOps & Cloud Professionals. Sign up for our newsletter for updated information, insight and promotion.