Put your best Weights and Biases model in production with UbiOps¶

Download notebook View source code

Weights and Biases (WandB) is a Machine Learning platform geared towards developers for building better models, faster. UbiOps is a Machine Learning Operations platform that allows you to put your best models into production quickly.

In this tutorial we show how to set-up a deployment in UbiOps that always uses the latest model that you selected in your WandB environment.

We first train a simple Convolutional Neural Network model, then we select the model with the highest accuracy, and place it in a dedicated artifact folder on WandB. Afterwards, we create and deploy a deployment in UbiOps that fetches the latest model from this artifact folder.

#This step may take a while
!pip install tensorflow==2.10.0
!pip install tensorflow_datasets==4.7.0
!pip install wandb==0.13.4
!pip install ubiops==3.15.0

To train a model on the MNIST dataset, we follow a publicly available TensorFlow Keras example.

First we prepare our dataset:

import tensorflow as tf
import tensorflow_datasets as tfds
#This step may also take a while
(ds_train, ds_test), ds_info = tfds.load(
    'mnist',
    split=['train', 'test'],
    shuffle_files=True,
    as_supervised=True,
    with_info=True,
)

def normalize_img(image, label):
  """Normalizes images: `uint8` -> `float32`."""
  return tf.cast(image, tf.float32) / 255., label

ds_train = ds_train.map(
    normalize_img, num_parallel_calls=tf.data.AUTOTUNE)
ds_train = ds_train.cache()
ds_train = ds_train.shuffle(ds_info.splits['train'].num_examples)
ds_train = ds_train.batch(128)
ds_train = ds_train.prefetch(tf.data.AUTOTUNE)

ds_test = ds_test.map(
    normalize_img, num_parallel_calls=tf.data.AUTOTUNE)
ds_test = ds_test.batch(128)
ds_test = ds_test.cache()
ds_test = ds_test.prefetch(tf.data.AUTOTUNE)

To log your training parameters to WandB, you first need to log-in to our WandB account. If this is your first time logging in, you are prompted to insert a valid WandB API token. Head to 'https://wandb.ai/settings' or 'https://wandb.ai/authorize' to retrieve the WandB API token for your account.

import wandb
wandb.login()

Train a Convolutional Neural Network model. Use the 'run' functionality from WandB to log the results on WandB. This step automatically creates a project named 'train-mnist-models' inside your WandB account.

from wandb.keras import WandbCallback
import os
#logging code

hyperparameters = dict(
    epochs = 7
)


with wandb.init(project="train-mnist-models", config=hyperparameters) as train_run:
    run_name = train_run.name
    config = train_run.config
    model = tf.keras.models.Sequential([
    tf.keras.layers.Flatten(input_shape=(28, 28)),
    tf.keras.layers.Dense(128, activation='relu'),
    tf.keras.layers.Dense(10)
    ])
    model.compile(
        optimizer=tf.keras.optimizers.Adam(0.001),
        loss=tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True),
        metrics=[tf.keras.metrics.SparseCategoricalAccuracy()],
    )

    #Wandbcallback integrates with model.fit. Make sure to enable 'save_model', so that wandb automatically saves 
    #the model with the highest accuracy.
    model.fit(ds_train,
        epochs=config.epochs,
        validation_data=ds_test,
        callbacks = [WandbCallback(save_model = True)]
    )

If we are happy with the results of this run, we can upload its best model to our dedicated artifact folder location.

# https://docs.wandb.ai/ref/python/public-api/artifact
with wandb.init(project = "production-environment", job_type = "model") as run:
    run_dir = train_run.dir
    artifact = wandb.Artifact('production-models', type = 'model')
    artifact.add_file(run_dir + r'\model-best.h5')
    run.log_artifact(artifact)

Put your model in production on UbiOps¶

We are now going to create a model that automatically queries the most recent model that has been placed inside the production environment. To do this, we initiate an empty deployment package folder and upload the required code to this folder.

!mkdir deployment_package

%%writefile deployment_package/requirements.txt

numpy == 1.23.4
imageio == 2.22.2
tensorflow  == 2.10.0
wandb == 0.13.4

%%writefile deployment_package/deployment.py

"""
The file containing the deployment code is required to be called 'deployment.py' and should contain the 'Deployment'
class and 'request' method. When a deployment container in UbiOps starts, the code in the __init__ method will run. This
is where we load the W&B artifacts. The code inside the request() method is run every time a call to the model API is
made.
"""

import os
from tensorflow.keras.models import load_model
from imageio import imread
import numpy as np
import wandb

class Deployment:

    def __init__(self, base_directory, context):

        print("Initialising deployment")
        # make connection to wand API here and pass reference to load_model
        wandb_api = wandb.Api()
        artifact_obj= wandb_api.artifact('production-environment/production-models:latest')
        artifact_obj.download('wandb-folder')
        model_path = os.path.join(base_directory, "wandb-folder/model-best.h5")
        self.model = load_model(model_path)

    def request(self, data):

        print("Processing request")

        x = imread(data['image'])
        # convert to a 4D tensor to feed into our model
        x = x.reshape(1, 28, 28, 1)
        x = x.astype(np.float32) / 255

        out = self.model.predict(x)
        prediction= int(np.argmax(out))

        # here we set our output parameters in the form of a json
        return {'prediction': prediction}

import shutil
shutil.make_archive('deployment_package', 'zip', '.', 'deployment_package')

Now that we have prepared our deployment package, we create a deployment on UbiOps, create a deployment version, and upload our deployment package to this deployment version.

import ubiops
API_TOKEN = 'Token abc123' # Fill in your token here
PROJECT_NAME = '<my-project-name>'    # Fill in your project name here
DEPLOYMENT_NAME = 'wandb-model'
DEPLOYMENT_VERSION = 'v1'


configuration = ubiops.Configuration(host="https://api.ubiops.com/v2.1")
configuration.api_key['Authorization'] = API_TOKEN

client = ubiops.ApiClient(configuration)
api = ubiops.CoreApi(client)
api.service_status()

mnist_template = ubiops.DeploymentCreate(
    name=DEPLOYMENT_NAME,
    description='A deployment to classify handwritten digits.',
    input_type='structured',
    output_type='structured',
    input_fields=[
        {'name': 'image', 'data_type': 'file'}
    ],
    output_fields=[
        {'name': 'prediction', 'data_type': 'string'}
    ]
)

mnist_deployment = api.deployments_create(project_name=PROJECT_NAME, data=mnist_template)
print(mnist_deployment)

Note that we need to add our WandB API key as an environment variable to our deployment, so that our deployment can connect to our WandB project.

WANDB_API_KEY = "<WandB API key>>"
api.deployment_environment_variables_create(
    PROJECT_NAME, 
    DEPLOYMENT_NAME, 
    data = {
        'name': "WANDB_API_KEY",  
        'secret': True, 
        'value': WANDB_API_KEY
    }
)

Create the deployment version

version_template = ubiops.DeploymentVersionCreate(
    version=DEPLOYMENT_VERSION,
    environment='python3-10',
    instance_type='1024mb',
    maximum_instances=1,
    minimum_instances=0,
    maximum_idle_time=1800, # = 30 minutes
    request_retention_mode='full',  # input/output of requests will be stored
    request_retention_time=3600  # requests will be stored for 1 hour
)

version = api.deployment_versions_create(
    project_name=PROJECT_NAME,
    deployment_name=DEPLOYMENT_NAME,
    data=version_template
)

And upload the deployment package

file_upload_result =api.revisions_file_upload(
    project_name=PROJECT_NAME,
    deployment_name=DEPLOYMENT_NAME,
    version=DEPLOYMENT_VERSION,
    file='deployment_package.zip'
)

Note that building a deployment can take long if this is the first time, because all packages from the requirements.txt need to be installed inside our deployment.

import time
ready = False
while not ready:
    response = api.deployment_versions_get(project_name=PROJECT_NAME,
        deployment_name=DEPLOYMENT_NAME, version = DEPLOYMENT_VERSION)
    if response.status == 'available':
        print("Deployment is ready!")
        break
    else:
        print("Deployment is not ready yet")
    time.sleep(15)

Now download an MNIST image depicting a "one", save it locally and push it to UbiOps as a file.

import requests
response = requests.get("https://storage.googleapis.com/ubiops/example-deployment-packages/1.jpg", stream = True)

with open('1.jpg', 'wb') as f:
        response.raw.decode_content = True
        shutil.copyfileobj(response.raw, f)

file_uri = ubiops.utils.upload_file(client, PROJECT_NAME, '1.jpg')
data = {'image': file_uri}

Lastly, send a request to the deployment.

api.deployment_requests_create(
    project_name=PROJECT_NAME,
    deployment_name=DEPLOYMENT_NAME,
    data=data
)

That's it! We have now created a set-up where we have an operational deployment that hosts the most recent version of your best model