How to access Artifacts for a Model Endpoint on Unified AI Platform when using Custom Containers for Prediction?

Viewed 479

Because of certain VPC restrictions I am forced to use custom containers for predictions for a model trained on Tensorflow. According to the documentation requirements I have created a HTTP server using Tensorflow Serving. The Dockerfile used to build the image is as follows:

FROM tensorflow/serving:2.3.0-gpu

# Set where models should be stored in the container
ENV MODEL_BASE_PATH=/models
RUN mkdir -p ${MODEL_BASE_PATH}

# copy the model file
ENV MODEL_NAME=my_model
COPY my_model /models/my_model

EXPOSE 5000

EXPOSE 8080

CMD ["tensorflow_model_server", "--rest_api_port=8080", "--port=5000", "--model_name=my_model", "--model_base_path=/models/my_model"]

Where my_model contains the saved_model inside a folder named 1/. I have then pushed the container image to Google Container Registry.

I would now like to pass a Model Artifact to this custom container such that I don't have to build and push a new docker image every time I train a new model. However I am unable to figure out how to access this new model (which is saved on a Cloud Storage Bucket) from within my Dockerfile while creating a Model on Unified AI Platform.

According to the documentation mentioned here the way to do so is as follows:

However, if you do provide model artifacts by specifying the artifactUri field, then the container must load these artifacts when it starts running. When AI Platform starts your container, it sets the AIP_STORAGE_URI environment variable to a Cloud Storage URI that begins with gs://. Your container's entrypoint command can download the directory specified by this URI in order to access the model artifacts.

However how do I rewrite the ENTRYPOINT to my Docker Image such that it reads the AIP_STORAGE_URI variable?

The link to the base image tensorflow/serving:2.3.0-gpu is here.

Any help will be appreciated.

1 Answers

In the documentation, it is explained that Vertex AI creates and manages a copy of the model artifacts that are passed when creating a model. The URI of the model artifact bucket managed by Vertex AI is stored in the AIP_STORAGE_URI environment variable. To access the model artifacts from this URI, the ENTRYPOINT command can be written as follows

ENTRYPOINT tensorflow_model_server --rest_api_port=8080 --port=5000 --model_name=$MODEL_NAME --model_base_path=$AIP_STORAGE_URI

The shell form of the ENTRYPOINT command is used here because there is a need to access the AIP_STORAGE_URI environment variable. More information about the usage of the ENTRYPOINT command can be found here.

Related