Deploy a TensorFlow model served with TensorFlow Serving using a custom container in an online endpoint

APPLIES TO: Azure CLI ml extension v2 (current)

APPLIES TO: Python SDK azure-ai-ml v2 (current)

Learn how to deploy a custom container as an online endpoint in Azure Machine Learning.

Custom container deployments can use web servers other than the default Python Flask server used by Azure Machine Learning. Users of these deployments can still take advantage of Azure Machine Learning's built-in monitoring, scaling, alerting, and authentication.

You can find various examples for TensorFlow Serving, TorchServe, Triton Inference Server, Plumber R package, and AzureML Inference Minimal image as below:

Example Script (CLI) Description
minimal/multimodel deploy-custom-container-minimal-multimodel Deploy multiple models to a single deployment by extending the AzureML Inference Minimal image.
minimal/single-model deploy-custom-container-minimal-single-model Deploy a single model by extending the AzureML Inference Minimal image.
mlflow/multideployment-scikit deploy-custom-container-mlflow-multideployment-scikit Deploy two MLFlow models with different Python requirements to two separate deployments behind a single endpoint using the AzureML Inference Minimal Image.
r/multimodel-plumber deploy-custom-container-r-multimodel-plumber Deploy three regression models to one endpoint using the Plumber R package
tfserving/half-plus-two deploy-custom-container-tfserving-half-plus-two Deploy a simple Half Plus Two model using a TensorFlow Serving custom container using the standard model registration process.
tfserving/half-plus-two-integrated deploy-custom-container-tfserving-half-plus-two-integrated Deploy a simple Half Plus Two model using a TensorFlow Serving custom container with the model integrated into the image.
torchserve/densenet deploy-custom-container-torchserve-densenet Deploy a single model using a TorchServe custom container.
torchserve/huggingface-textgen deploy-custom-container-torchserve-huggingface-textgen Deploy Hugging Face models to an online endpoint and follow along with the Hugging Face Transformers TorchServe example.
triton/single-model deploy-custom-container-triton-single-model Deploy a Triton model using a custom container

This article focuses on serving a TensorFlow model with TensorFlow (TF) Serving.

Warning

Microsoft may not be able to help troubleshoot problems caused by a custom image. If you encounter problems, you may be asked to use the default image or one of the images Microsoft provides to see if the problem is specific to your image.

Prerequisites

Before following the steps in this article, make sure you have the following prerequisites:

  • You, or the service principal you use, must have Contributor access to the Azure Resource Group that contains your workspace. You'll have such a resource group if you configured your workspace using the quickstart article.

  • To deploy locally, you must have Docker engine running locally. This step is highly recommended. It will help you debug issues.

Download source code

To follow along with this tutorial, download the source code below.

git clone https://github.com/Azure/azureml-examples --depth 1
cd azureml-examples/cli

Initialize environment variables

Define environment variables:

BASE_PATH=endpoints/online/custom-container/tfserving/half-plus-two
AML_MODEL_NAME=tfserving-mounted
MODEL_NAME=half_plus_two
MODEL_BASE_PATH=/var/azureml-app/azureml-models/$AML_MODEL_NAME/1

Download a TensorFlow model

Download and unzip a model that divides an input by two and adds 2 to the result:

wget https://aka.ms/half_plus_two-model -O $BASE_PATH/half_plus_two.tar.gz
tar -xvf $BASE_PATH/half_plus_two.tar.gz -C $BASE_PATH

Run a TF Serving image locally to test that it works

Use docker to run your image locally for testing:

docker run --rm -d -v $PWD/$BASE_PATH:$MODEL_BASE_PATH -p 8501:8501 \
 -e MODEL_BASE_PATH=$MODEL_BASE_PATH -e MODEL_NAME=$MODEL_NAME \
 --name="tfserving-test" docker.io/tensorflow/serving:latest
sleep 10

Check that you can send liveness and scoring requests to the image

First, check that the container is "alive," meaning that the process inside the container is still running. You should get a 200 (OK) response.

curl -v http://localhost:8501/v1/models/$MODEL_NAME

Then, check that you can get predictions about unlabeled data:

curl --header "Content-Type: application/json" \
  --request POST \
  --data @$BASE_PATH/sample_request.json \
  http://localhost:8501/v1/models/$MODEL_NAME:predict

Stop the image

Now that you've tested locally, stop the image:

docker stop tfserving-test

Deploy your online endpoint to Azure

Next, deploy your online endpoint to Azure.

Create a YAML file for your endpoint and deployment

You can configure your cloud deployment using YAML. Take a look at the sample YAML for this example:

tfserving-endpoint.yml

$schema: https://azuremlsdk2.blob.core.windows.net/latest/managedOnlineEndpoint.schema.json
name: tfserving-endpoint
auth_mode: aml_token

tfserving-deployment.yml

$schema: https://azuremlschemas.azureedge.net/latest/managedOnlineDeployment.schema.json
name: tfserving-deployment
endpoint_name: tfserving-endpoint
model:
  name: tfserving-mounted
  version: {{MODEL_VERSION}}
  path: ./half_plus_two
environment_variables:
  MODEL_BASE_PATH: /var/azureml-app/azureml-models/tfserving-mounted/{{MODEL_VERSION}}
  MODEL_NAME: half_plus_two
environment:
  #name: tfserving
  #version: 1
  image: docker.io/tensorflow/serving:latest
  inference_config:
    liveness_route:
      port: 8501
      path: /v1/models/half_plus_two
    readiness_route:
      port: 8501
      path: /v1/models/half_plus_two
    scoring_route:
      port: 8501
      path: /v1/models/half_plus_two:predict
instance_type: Standard_DS2_v2
instance_count: 1

There are a few important concepts to notice in this YAML/Python parameter:

Readiness route vs. liveness route

An HTTP server defines paths for both liveness and readiness. A liveness route is used to check whether the server is running. A readiness route is used to check whether the server is ready to do work. In machine learning inference, a server could respond 200 OK to a liveness request before loading a model. The server could respond 200 OK to a readiness request only after the model has been loaded into memory.

Review the Kubernetes documentation for more information about liveness and readiness probes.

Notice that this deployment uses the same path for both liveness and readiness, since TF Serving only defines a liveness route.

Locating the mounted model

When you deploy a model as an online endpoint, Azure Machine Learning mounts your model to your endpoint. Model mounting enables you to deploy new versions of the model without having to create a new Docker image. By default, a model registered with the name foo and version 1 would be located at the following path inside of your deployed container: /var/azureml-app/azureml-models/foo/1

For example, if you have a directory structure of /azureml-examples/cli/endpoints/online/custom-container on your local machine, where the model is named half_plus_two:

Diagram showing a tree view of the local directory structure.

and tfserving-deployment.yml contains:

model:
    name: tfserving-mounted
    version: 1
    path: ./half_plus_two

then your model will be located under /var/azureml-app/azureml-models/tfserving-deployment/1 in your deployment:

Diagram showing a tree view of the deployment directory structure.

You can optionally configure your model_mount_path. It enables you to change the path where the model is mounted.

Important

The model_mount_path must be a valid absolute path in Linux (the OS of the container image).

For example, you can have model_mount_path parameter in your tfserving-deployment.yml:

name: tfserving-deployment
endpoint_name: tfserving-endpoint
model:
  name: tfserving-mounted
  version: 1
  path: ./half_plus_two
model_mount_path: /var/tfserving-model-mount
.....

then your model will be located at /var/tfserving-model-mount/tfserving-deployment/1 in your deployment. Note that it's no longer under azureml-app/azureml-models, but under the mount path you specified:

Diagram showing a tree view of the deployment directory structure when using mount_model_path.

Create your endpoint and deployment

Now that you've understood how the YAML was constructed, create your endpoint.

az ml online-endpoint create --name tfserving-endpoint -f endpoints/online/custom-container/tfserving-endpoint.yml

Creating a deployment may take few minutes.

az ml online-deployment create --name tfserving-deployment -f endpoints/online/custom-container/tfserving-deployment.yml --all-traffic

Invoke the endpoint

Once your deployment completes, see if you can make a scoring request to the deployed endpoint.

RESPONSE=$(az ml online-endpoint invoke -n $ENDPOINT_NAME --request-file $BASE_PATH/sample_request.json)

Delete the endpoint

Now that you've successfully scored with your endpoint, you can delete it:

az ml online-endpoint delete --name tfserving-endpoint

Next steps