Model Registry Service with Wallaroo Tutorials

TensorFlow File Format

TensorFlow models are .zip file of the SavedModel format. For example, the Aloha sample TensorFlow model is stored in the directory alohacnnlstm:

├── saved_model.pb
└── variables
    ├── variables.data-00000-of-00002
    ├── variables.data-00001-of-00002
    └── variables.index

This is compressed into the .zip file alohacnnlstm.zip with the following command:

zip -r alohacnnlstm.zip alohacnnlstm/

ML models that meet the Tensorflow and SavedModel format will run as Wallaroo Native runtimes by default.

See the SavedModel guide for full details.

Python models uploaded to Wallaroo are executed Wallaroo Containerized Runtime.

Note that Python models - aka “Python steps” - are standalone python scripts that use the python libraries. These are commonly used for data formatting such as the pre and post-processing steps, and are also appropriate for simple models (such as ARIMA Statsmodels). A Wallaroo Python model can be composed of one or more Python script that matches the Wallaroo requirements.

This is contrasted with Arbitrary Python models, also known as Bring Your Own Predict (BYOP) that allow for custom model inference methods with supporting scripts and artifacts. These are used with pre-trained models (PyTorch, Tensorflow, etc) along with their supporting artifacts such as other Python modules, scripts, model files, etc.

Python Models Requirements

Python scripts packaged as Python models in Wallaroo have the following requirements.

• At least one .py Python script file with the following:

  • Must be compatible with Python version 3.8.

  • Imports the mac.types.InferenceData included with the Wallaroo SDK. For example:

    from mac.types import InferenceData
    

  • Includes the following method as the entry point for Wallaroo model inferencing:

    def process_data(input_data: InferenceData) -> InferenceData:
        # additional code block here
    

    • Only one implementation of process_data(input_data: InferenceData) -> InferenceData is allowed. There can be as many Python scripts included in the .zip file as needed, but only one can have this method as the entry point.

      • The process_data function must return a dictionary where the keys are strings and the values are NumPy arrays. In the case of single values (scalars) these must be single-element arrays. For example:

        def process_data(input_data: InferenceData) -> InferenceData:
         # return a dictionary with the field output that transforms the input field `variable` to its value to the 10th power.
          return {
            'output' : np.rint(np.power(10, input_data["variable"]))
          }
        

    • InferenceData represents a dictionary of numpy arrays where the first dimension is always the batch size. The type annotations set in the input_schema and output_schema for the model when uploaded must be present and correct.

    • process_data accepts and returns InferenceData. Any other implementations will return an error.

  • (Optional): A requirements.txt file that includes any additional Python libraries required by the Python script with the following requirements:

    • The Python libraries must match the targeted infrastructure. For details on uploading models to a specific infrastructures such as ARM, see Automated Model Packaging.
    • The Python libraries must be compatible with Python version python==3.8.6.

The Python script, optional requirements.txt file, and artifacts are packaged in a .zip file with the Python script and optional requirements.txt file in the root folder. For example, the sample files stored in the folder preprocess-step:

/preprocss-step
    sample-script.py
    requirements.txt
    /artifacts
        datalist.csv

The files are packaged into a .zip file. For example, the following packages the contents of the folder preprocess_step into preprocess_step.zip.

zip -r preprocess_step.zip preprocess_step/*

In the example below, the Python model is used as a pre processing step for another ML model. It accepts as an input the InferenceData submitted as part of an inference request. It then formats the data and outputs a dictionary of numpy arrays with the field tensor. This data is then able to be passed to the next model in a pipeline step.

import datetime
import logging

import numpy as np
import pandas as pd

import wallaroo

from mac.types import InferenceData

logger = logging.getLogger(__name__)

_vars = [
    "bedrooms",
    "bathrooms",
    "sqft_living",
    "sqft_lot",
    "floors",
    "waterfront",
    "view",
    "condition",
    "grade",
    "sqft_above",
    "sqft_basement",
    "lat",
    "long",
    "sqft_living15",
    "sqft_lot15",
    "house_age",
    "renovated",
    "yrs_since_reno",
]


def process_data(input_data: InferenceData) -> InferenceData:
    input_df = pd.DataFrame(input_data)
    thisyear = datetime.datetime.now().year
    input_df["house_age"] = thisyear - input_df["yr_built"]
    input_df["renovated"] = np.where((input_df["yr_renovated"] > 0), 1, 0)
    input_df["yrs_since_reno"] = np.where(
        input_df["renovated"],
        input_df["yr_renovated"] - input_df["yr_built"],
        0,
    )
    input_df = input_df.loc[:, _vars]

    return {"tensor": input_df.to_numpy(dtype=np.float32)}

In line with other Wallaroo inference results, the outputs of a Python step that returns a pandas DataFrame or Arrow Table will be listed in the out. metadata, with all inference outputs listed as out.{variable 1}, out.{variable 2}, etc. For example, a postprocessing Python step that is the final model step in a pipeline with the output field output is included in the out dataset as the field out.output in the Wallaroo inference result.

During the model upload process, Wallaroo optimizes models by converting them to the Wallaroo Native Runtime, if possible, or running the model directly in the Wallaroo Containerized Runtime. See the Model Deploy for details on how to configure pipeline resources based on the model’s runtime.

• IMPORTANT CONFIGURATION NOTE: For PyTorch input schemas, the floats must be pyarrow.float32() for the PyTorch model to be converted to the Native Wallaroo Runtime during the upload process.

PyTorch Input and Output Schemas

PyTorch input and output schemas have additional requirements depending on whether the PyTorch model is single input/output or multiple input/output. This refers to the number of columns:

• Single Input/Output: Has one input and one output column.
• Multiple Input/Output: Has more than one input or more than one output column.

The column names for the model can be anything. For example:

• Model Input Fields:
  • length
  • width
  • intensity
  • etc

When creating the input and output schemas for uploading a PyTorch model in Wallaroo, the field names must match the following requirements. For example, for multi-column PyTorch models, the input would be:

• Data Schema Input Fields:
  • input_1
  • input_2
  • input_3
  • input_...

For single input/output PyTorch model, the field names must be input and output. For example, if the input field is a List of Floats of size 10, and the output field is a list of floats of list size one, the input and output schemas are:

input_schema = pa.schema([
    pa.field('input', pa.list_(pa.float32(), list_size=10))
])

output_schema = pa.schema([
    pa.field('output', pa.list_(pa.float32(), list_size=1))
])

For multi input/output PyTorch models, the data schemas for each input and output field must be named input_1, input_2... and output_1, output_2, etc. These must be in the same order that the PyTorch model is trained to accept them.

For example, a multi input/output PyTorch model that takes the following inputs and outputs:

• Inputs
  • input_1: List of Floats of length 10.
  • input_2: List of Floats of length 5.
• Outputs
  • output_1: List of Floats of length 3.
  • output_2: List of Floats of length 2.

The following input and output schemas would be used.

input_schema = pa.schema([
    pa.field('input_1', pa.list_(pa.float32(), list_size=10)),
    pa.field('input_2', pa.list_(pa.float32(), list_size=5))
])
output_schema = pa.schema([
    pa.field('output_1', pa.list_(pa.float32(), list_size=3)),
    pa.field('output_2', pa.list_(pa.float32(), list_size=2))
])

Sci-kit Learn aka SKLearn.

During the model upload process, Wallaroo optimizes models by converting them to the Wallaroo Native Runtime, if possible, or running the model directly in the Wallaroo Containerized Runtime. See the Model Deploy for details on how to configure pipeline resources based on the model’s runtime.

SKLearn Schema Inputs

SKLearn schema follows a different format than other models. To prevent inputs from being out of order, the inputs should be submitted in a single row in the order the model is trained to accept, with all of the data types being the same. For example, the following DataFrame has 4 columns, each column a float.

For submission to an SKLearn model, the data input schema will be a single array with 4 float values.

input_schema = pa.schema([
    pa.field('inputs', pa.list_(pa.float64(), list_size=4))
])

When submitting as an inference, the DataFrame is converted to rows with the column data expressed as a single array. The data must be in the same order as the model expects, which is why the data is submitted as a single array rather than JSON labeled columns: this insures that the data is submitted in the exact order as the model is trained to accept.

Original DataFrame:

Converted DataFrame:

SKLearn Schema Outputs

Outputs for SKLearn that are meant to be predictions or probabilities when output by the model are labeled in the output schema for the model when uploaded to Wallaroo. For example, a model that outputs either 1 or 0 as its output would have the output schema as follows:

output_schema = pa.schema([
    pa.field('predictions', pa.int32())
])

When used in Wallaroo, the inference result is contained in the out metadata as out.predictions.

pipeline.infer(dataframe)

During the model upload process, Wallaroo optimizes models by converting them to the Wallaroo Native Runtime, if possible, or running the model directly in the Wallaroo Containerized Runtime. See the Model Deploy for details on how to configure pipeline resources based on the model’s runtime.

TensorFlow Keras SavedModel Format

TensorFlow Keras SavedModel models are .zip file of the SavedModel format. For example, the Aloha sample TensorFlow model is stored in the directory alohacnnlstm:

├── saved_model.pb
└── variables
    ├── variables.data-00000-of-00002
    ├── variables.data-00001-of-00002
    └── variables.index

This is compressed into the .zip file alohacnnlstm.zip with the following command:

zip -r alohacnnlstm.zip alohacnnlstm/

See the SavedModel guide for full details.

TensorFlow Keras H5 Format

Wallaroo supports the H5 for Tensorflow Keras models.

During the model upload process, Wallaroo optimizes models by converting them to the Wallaroo Native Runtime, if possible, or running the model directly in the Wallaroo Containerized Runtime. See the Model Deploy for details on how to configure pipeline resources based on the model’s runtime.

Since the Wallaroo 2024.1 release, XGBoost support is enhanced to performantly support a wider set of XGBoost models. XGBoost models are not required to be trained with ONNX nomenclature in order to successfully convert to a performant runtime.

XGBoost Types Support

The following XGBoost model types are supported by Wallaroo. XGBoost models not supported by Wallaroo are supported via the Arbitrary Python models, also known as Bring Your Own Predict (BYOP).

• XGBRanker XGBoost models are currently supported via converting them to BYOP models.

XGBoost Schema Inputs

XGBoost schema follows a different format than other models. To prevent inputs from being out of order, the inputs should be submitted in a single row in the order the model is trained to accept, with all of the data types being the same. If a model is originally trained to accept inputs of different data types, it will need to be retrained to only accept one data type for each column - typically pa.float64() is a good choice.

For example, the following DataFrame has 4 columns, each column a float.

For submission to an XGBoost model, the data input schema will be a single array with 4 float values.

input_schema = pa.schema([
    pa.field('inputs', pa.list_(pa.float64(), list_size=4))
])

When submitting as an inference, the DataFrame is converted to rows with the column data expressed as a single array. The data must be in the same order as the model expects, which is why the data is submitted as a single array rather than JSON labeled columns: this insures that the data is submitted in the exact order as the model is trained to accept.

Original DataFrame:

Converted DataFrame:

XGBoost Schema Outputs

Outputs for XGBoost are labeled based on the trained model outputs.

Outputs for XBoost that are meant to be predictions or probabilities must be labeled as part of the output schema. For example, a model that outputs either 1 or 0 as its output would have the output schema as follows:

output_schema = pa.schema([
    pa.field('predictions', pa.int32())
])

When used in Wallaroo, the inference result is contained in the out metadata as out.predictions.

pipeline.infer(dataframe)

Arbitrary Python models, also known as Bring Your Own Predict (BYOP) allow for custom model inference methods with supporting scripts and artifacts. These are used with pre-trained models (PyTorch, Tensorflow, etc) along with their supporting artifacts such as other Python modules, scripts, model files, etc.

Contrast this with Wallaroo Python models - aka “Python steps” - are standalone python scripts that use the python libraries. These are commonly used for data formatting such as the pre and post-processing steps, and are also appropriate for simple models (such as ARIMA Statsmodels). A Wallaroo Python model can be composed of one or more Python script that matches the Wallaroo requirements.

Arbitrary Python File Requirements

Arbitrary Python (BYOP) models are uploaded to Wallaroo via a ZIP file with the following components:

For example, the if the arbitrary python model will be known as vgg_clustering, the contents may be in the following structure, with vgg_clustering as the storage directory:

vgg_clustering\
    feature_extractor.h5
    kmeans.pkl
    custom_inference.py
    requirements.txt

Note the inclusion of the custom_inference.py file. This file name is not required - any Python script or scripts that extend the classes listed above are sufficient. This Python script could have been named vgg_custom_model.py or any other name as long as it includes the extension of the classes listed above.

The sample arbitrary python model file is created with the command zip -r vgg_clustering.zip vgg_clustering/.

Wallaroo Arbitrary Python uses the Wallaroo SDK mac module, included in the Wallaroo SDK 2023.2.1 and above. See the Wallaroo SDK Install Guides for instructions on installing the Wallaroo SDK.

Arbitrary Python Script Requirements

The entry point of the arbitrary python model is any python script that extends the following classes. These are included with the Wallaroo SDK. The required methods that must be overridden are specified in each section below.

• mac.inference.Inference interface serves model inferences based on submitted input some input. Its purpose is to serve inferences for any supported arbitrary model framework (e.g. scikit, keras etc.).

  classDiagram
      class Inference {
          <<Abstract>>
          +model Optional[Any]
          +expected_model_types()* Set
          +predict(input_data: InferenceData)*  InferenceData
          -raise_error_if_model_is_not_assigned() None
          -raise_error_if_model_is_wrong_type() None
      }

• mac.inference.creation.InferenceBuilder builds a concrete Inference, i.e. instantiates an Inference object, loads the appropriate model and assigns the model to to the Inference object.

  classDiagram
      class InferenceBuilder {
          +create(config InferenceConfig) * Inference
          -inference()* Any
      }

mac.inference.Inference

mac.inference.Inference Objects

mac.inference.Inference Methods

preds = self.model.predict(data)
preds = preds.numpy()
rows, _ = preds.shape
preds = preds.reshape((rows,))

return {"prediction": preds} # a Dictionary of numpy arrays.

The example, the expected_model_types can be defined for the KMeans model.

from sklearn.cluster import KMeans

class SampleClass(mac.inference.Inference):
    @property
    def expected_model_types(self) -> Set[Any]:
        return {KMeans}

mac.inference.creation.InferenceBuilder

InferenceBuilder builds a concrete Inference, i.e. instantiates an Inference object, loads the appropriate model and assigns the model to the Inference.

classDiagram
    class InferenceBuilder {
        +create(config InferenceConfig) * Inference
        -inference()* Any
    }

Each model that is included requires its own InferenceBuilder. InferenceBuilder loads one model, then submits it to the Inference class when created. The Inference class checks this class against its expected_model_types() Set.

mac.inference.creation.InferenceBuilder Methods

Arbitrary Python Runtime

Arbitrary Python always run in the containerized model runtime.

Arbitrary Python Inputs

Arbitrary Python inputs are defined during model upload in Apache Arrow Schema format with the following conditions:

• By default, data inputs are optional unless they are specified with nullable=False.
• The arbitrary Python code must be aware of the optional and required fields and how to manage those inputs.
• Specific Data Types conditions:
  • Scaler: Scaler values can be Null.
  • Lists: Lists must either be empty [] or an an array of Null values, for example [None], but cannot be passed as Null outside of an array.
• By default, columns with only the None or Null value are assigned by Python as NullArray, which is an array with all values of Null. In these situations, the schema must be specified.

import wallaroo
import pyarrow as pa
input_schema = pa.schema([
    pa.field('input_1', pa.list_(pa.float32()), nullable=False), # fields are optional by default unless `nullable` is set to `False`
    pa.field('input_2', pa.list_(pa.float32())),
    pa.field('multiply_factor', pa.int32()),
])

output_schema = pa.schema([
    pa.field('output', pa.list_(pa.float32())),
])

input_1 = [[1., 2.], [3., 4.]]
input_2 = [[5., 6.], [7., 8.]]
multiply_factor = [2, 3]
arrow_table = pa.table({"input_1": input_1, "input_2": input_2, "multiply_factor": multiply_factor}, schema=input_schema)
display(arrow_)table

input_1 = [[1., 2.], [3., 4.]]
input_2 = [[], []]
multiply_factor = [None, None]
arrow_table = pa.table({"input_1": input_1, "input_2": input_2, "multiply_factor": multiply_factor}, schema=input_schema)
arrow_table

pipeline.infer(arrow_table)

pyarrow.Table
time: timestamp[ms]
in.input_1: list<item: float> not null
  child 0, item: float
in.input_2: list<item: float> not null
  child 0, item: float
in.multiply_factor: int32 not null
out.output: list<item: double> not null
  child 0, item: double
anomaly.count: uint32 not null
----
time: [[2024-04-30 09:12:01.445,2024-04-30 09:12:01.445]]
in.input_1: [[[1,2],[3,4]]]
in.input_2: [[[5,6],[7,8]]]
in.multiply_factor: [[2,3]]
out.output: [[[12,16],[30,36]]]
anomaly.count: [[0,0]]

dataframe = pd.DataFrame({'input_1': [[1., 2.], [3., 4.]], 'input_2': [[], []], 'multiply_factor': [2, 3]})
display(dataframe)

input_1 = [[1., 2.], [3., 4.]]
input_2 = [[], []]
multiply_factor = [None, None]
arrow_table = pa.table({"input_1": input_1, "input_2": input_2, "multiply_factor": multiply_factor}, schema=input_schema)
display(arrow_table)

pyarrow.Table
input_1: list<item: float> not null
  child 0, item: float
input_2: list<item: float>
  child 0, item: float
multiply_factor: int32
----
input_1: [[[1,2],[3,4]]]
input_2: [[[],[]]]
multiply_factor: [[null,null]]

pipeline.infer(arrow_table)

pyarrow.Table
time: timestamp[ms]
in.input_1: list<item: float> not null
  child 0, item: float
in.input_2: list<item: float> not null
  child 0, item: float
in.multiply_factor: int32 not null
out.output: list<item: double> not null
  child 0, item: double
anomaly.count: uint32 not null
----
time: [[2024-04-30 09:07:42.467,2024-04-30 09:07:42.467]]
in.input_1: [[[1,2],[3,4]]]
in.input_2: [[[],[]]]
in.multiply_factor: [[null,null]]
out.output: [[[1,2],[3,4]]]
anomaly.count: [[0,0]]

For models that do not fall under the supported model frameworks, organizations can use containerized MLFlow ML Models.

This guide details how to add ML Models from a model registry service into Wallaroo.

Wallaroo supports both public and private containerized model registries. See the Wallaroo Private Containerized Model Container Registry Guide for details on how to configure a Wallaroo instance with a private model registry.

Wallaroo users can register their trained MLFlow ML Models from a containerized model container registry into their Wallaroo instance and perform inferences with it through a Wallaroo pipeline.

As of this time, Wallaroo only supports MLFlow 1.30.0 containerized models. For information on how to containerize an MLFlow model, see the MLFlow Documentation.