Deploy a Self-Hosted Model Registry

Using S3 as a simple registry to store your ML/DL Models

Photo by Valeriia Svitlini on Unsplash

Model registries allow Data Scientists to store, track, and load their models for continual experimentation and eventual deployment. There are tons of model registries out there, but what if you want to host your own on your AWS account?

Using the model-registry repo from my GitHub, you can easily deploy a registry in an S3 bucket and upload your first model in minutes. It’s easy and lightweight, and the costs are only what you incur from your S3 bucket. Let’s take a look.

Using S3

S3 is usually the first AWS resource that people interact with. It’s simple to use, easy to understand, and cheap. For our lightweight model registry, it’s great. It’s essentially a NoSQL database that we can write files to in any format. It can store simple files like JSON but also more complex files like our ML/DL models.

To set up an S3 bucket as our model registry, simply clone the repo and follow the setup instructions in the README. Namely, add a .env file to store your AWS variables and registry name.

.env

AWS_ACCOUNT_ID=1234567890
AWS_DEFAULT_REGION=ap-southeast-1
REGISTRY_NAME=my-model-registry

Next, run the following Makefile command. This will pick up your variables from the .env file and pass them to Terraform. Terraform will then deploy your S3 bucket to your AWS account.

make deploy-registry

When this is complete, it will print out the name of the bucket as bucket_name in your terminal. Export that to your environment using the following command:

export REGISTRY_BUCKET=<bucket_name_from_terraform_output>

Your model registry is deployed and you’re ready to upload your first model!

Install the model-registry library

To interact with your registry, you’ll need to install the model-registry library in the repo. You can do this easily by running:

make install

This will install model-registry in your current Python environment. You can now use the library anywhere you are using this environment including Jupyter notebooks, Python scripts, etc.

Creating, Saving, Loading, and Deleting Models

Now you’re ready to interact with your registry!

If you’d like to skip the step-by-step and see how the library works end-to-end, you can run some of the scripts in the examples/ directory. Just call make run-examples.

Uploading a Model

First, let’s create a deep-learning network. We’ll use PyTorch as an example.

import torch 

class NeuralNetwork(torch.nn.Module):
    def __init__(self):
        super().__init__()
        self.flatten = torch.nn.Flatten()
        self.linear_relu_stack = torch.nn.Sequential(
            torch.nn.Linear(28 * 28, 512),
            torch.nn.ReLU(),
            torch.nn.Linear(512, 512),
            torch.nn.ReLU(),
            torch.nn.Linear(512, 10),
        )

    def forward(self, x):
        x = self.flatten(x)
        logits = self.linear_relu_stack(x)
        return logits

Now, to upload this model we need to wrap it in one of the model-registry library’s Model classes. Let’s import the PyTorch model and wrap our network in it.

from registry.torch import TorchModel
from registry.registry import ModelRegistry

# create model
model = TorchModel(obj=NeuralNetwork(), name="my_torch_model")

Notice we need to also provide a name. This is simply to have a human-friendly alias to reference.

Now that our model is created, we can upload it to our registry. We’ll import the ModelRegistry class, instantiate it, and upload our model. Note that during instantiation the ModelRegistry will look for the REGISTRY_BUCKET environment variable that we set. If the variable is not there, an error will be thrown.

from registry.registry import ModelRegistry

# upload to registry
registry = ModelRegistry()
registry.upload_model(model)

Our model and it’s metadata has been uploaded to our registry! You can verify it’s there by going to S3 in your AWS console.

Loading Models

Loading your models from your registry is just as simple. We just need to pass a model_id to the load_model method. This value is available in each model object.

# get model_id
model_id = model.id

# load model from registry
registry = ModelRegistry()
loaded_model = registry.load_model(model_id)

Alternatively, you can list out all of your model_id s in the registry and load one from the list.

# list model_ids
registry = ModelRegistry()
model_ids = registry.list_models()
print(model_ids)

loaded_model = registry.load_model(model_ids[0])

Model Versioning

Each subsequent time you upload a model object, the version will be auto-incremented. This is useful in tracking changes to the same model.

# create model
model = TorchModel(obj=NeuralNetwork(), name="my_torch_model")

# first upload = version 1
loaded_model = registry.load_model(model_ids[0])
print(loaded_model.version)
# > 1

# perform some changes to the model...

# upload again (version 2)
registry.upload_model(loaded_model)
print(loaded_model.version)
# > 2

You can check all the versions of a model by passing its ID to the registry.

# get all versions of our model
versions = registry.list_model_versions(model.id)
print(versions)
# > [1,2]

Many of the methods on ModelRegistry have an optional model_version parameter. If this is left blank, the newest version is implied.

Deleting a Model

Finally, you can clean up your registry by deleting models when you are sure you no longer need them.

# delete our model
registry.delete_model(model_id)

Conclusion

This GitHub repo helps you set up a lightweight, easy-to-use, and cheap model registry for your Data Science / Machine Learning / Deep Learning projects. In as little as three commands you can have your own self-hosted model registry ready to use.

The code for this post can be found on my GitHub page.

If you liked this post or have any questions/comments, please leave a comment below!

Want updates on what I’m writing, side projects I’m building, or articles I find interesting? Feel free to follow me.

Thanks for reading! 👨🏼‍💻