Wallaroo users can register their trained machine learning models from a model registry into their Wallaroo instance and perform inferences with it through a Wallaroo pipeline.

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

Artifact Requirements

Models are uploaded to the Wallaroo instance as the specific artifact - the “file” or other data that represents the file itself. This must comply with the Wallaroo model requirements framework and version or it will not be deployed. Note that for models that fall outside of the supported model types, they can be registered to a Wallaroo workspace as MLFlow 1.30.0 containerized models.

Supported Models

The following frameworks are supported. Frameworks fall under either Native or Containerized runtimes in the Wallaroo engine. For more details, see the specific framework what runtime a specific model framework runs in.

Please note the following.

• Native Runtime Frameworks:
• ONNX
• Tensorflow
• Python

Wallaroo natively supports Open Neural Network Exchange (ONNX) models into the Wallaroo engine.

Parameter	Description
Web Site	https://onnx.ai/
Supported Libraries	See table below.
Framework	Framework.ONNX aka onnx
Runtime	Native aka onnx

The following ONNX versions models are supported:

Wallaroo Version	ONNX Version	ONNX IR Version	ONNX OPset Version	ONNX ML Opset Version
2023.2.1 (July 2023)	1.12.1	8	17	3
2023.2 (May 2023)	1.12.1	8	17	3
2023.1 (March 2023)	1.12.1	8	17	3
2022.4 (December 2022)	1.12.1	8	17	3
After April 2022 until release 2022.4 (December 2022)	1.10.*	7	15	2
Before April 2022	1.6.*	7	13	2

For the most recent release of Wallaroo 2023.2.1, the following native runtimes are supported:

• If converting another ML Model to ONNX (PyTorch, XGBoost, etc) using the onnxconverter-common library, the supported DEFAULT_OPSET_NUMBER is 17.

Using different versions or settings outside of these specifications may result in inference issues and other unexpected behavior.

ONNX models always run in the native runtime space.

Data Schemas

ONNX models deployed to Wallaroo have the following data requirements.

• Equal rows constraint: The number of input rows and output rows must match.
• All inputs are tensors: The inputs are tensor arrays with the same shape.
• Data Type Consistency: Data types within each tensor are of the same type.

Equal Rows Constraint

Inference performed through ONNX models are assumed to be in batch format, where each input row corresponds to an output row. This is reflected in the in fields returned for an inference. In the following example, each input row for an inference is related directly to the inference output.

df = pd.read_json('./data/cc_data_1k.df.json')
display(df.head())
result = ccfraud_pipeline.infer(df.head())

display(result)

INPUT

	tensor
0	[-1.0603297501, 2.3544967095000002, -3.5638788326, 5.1387348926, -1.2308457019, -0.7687824608, -3.5881228109, 1.8880837663, -3.2789674274, -3.9563254554, 4.0993439118, -5.6539176395, -0.8775733373, -9.131571192000001, -0.6093537873, -3.7480276773, -5.0309125017, -0.8748149526000001, 1.9870535692, 0.7005485718000001, 0.9204422758, -0.1041491809, 0.3229564351, -0.7418141657, 0.0384120159, 1.0993439146, 1.2603409756, -0.1466244739, -1.4463212439]
1	[-1.0603297501, 2.3544967095000002, -3.5638788326, 5.1387348926, -1.2308457019, -0.7687824608, -3.5881228109, 1.8880837663, -3.2789674274, -3.9563254554, 4.0993439118, -5.6539176395, -0.8775733373, -9.131571192000001, -0.6093537873, -3.7480276773, -5.0309125017, -0.8748149526000001, 1.9870535692, 0.7005485718000001, 0.9204422758, -0.1041491809, 0.3229564351, -0.7418141657, 0.0384120159, 1.0993439146, 1.2603409756, -0.1466244739, -1.4463212439]
2	[-1.0603297501, 2.3544967095000002, -3.5638788326, 5.1387348926, -1.2308457019, -0.7687824608, -3.5881228109, 1.8880837663, -3.2789674274, -3.9563254554, 4.0993439118, -5.6539176395, -0.8775733373, -9.131571192000001, -0.6093537873, -3.7480276773, -5.0309125017, -0.8748149526000001, 1.9870535692, 0.7005485718000001, 0.9204422758, -0.1041491809, 0.3229564351, -0.7418141657, 0.0384120159, 1.0993439146, 1.2603409756, -0.1466244739, -1.4463212439]
3	[-1.0603297501, 2.3544967095000002, -3.5638788326, 5.1387348926, -1.2308457019, -0.7687824608, -3.5881228109, 1.8880837663, -3.2789674274, -3.9563254554, 4.0993439118, -5.6539176395, -0.8775733373, -9.131571192000001, -0.6093537873, -3.7480276773, -5.0309125017, -0.8748149526000001, 1.9870535692, 0.7005485718000001, 0.9204422758, -0.1041491809, 0.3229564351, -0.7418141657, 0.0384120159, 1.0993439146, 1.2603409756, -0.1466244739, -1.4463212439]
4	[0.5817662108, 0.09788155100000001, 0.1546819424, 0.4754101949, -0.19788623060000002, -0.45043448540000003, 0.016654044700000002, -0.0256070551, 0.0920561602, -0.2783917153, 0.059329944100000004, -0.0196585416, -0.4225083157, -0.12175388770000001, 1.5473094894000001, 0.2391622864, 0.3553974881, -0.7685165301, -0.7000849355000001, -0.1190043285, -0.3450517133, -1.1065114108, 0.2523411195, 0.0209441826, 0.2199267436, 0.2540689265, -0.0450225094, 0.10867738980000001, 0.2547179311]

OUTPUT

	time	in.tensor	out.dense_1	check_failures
0	2023-11-17 20:34:17.005	[-1.0603297501, 2.3544967095, -3.5638788326, 5.1387348926, -1.2308457019, -0.7687824608, -3.5881228109, 1.8880837663, -3.2789674274, -3.9563254554, 4.0993439118, -5.6539176395, -0.8775733373, -9.131571192, -0.6093537873, -3.7480276773, -5.0309125017, -0.8748149526, 1.9870535692, 0.7005485718, 0.9204422758, -0.1041491809, 0.3229564351, -0.7418141657, 0.0384120159, 1.0993439146, 1.2603409756, -0.1466244739, -1.4463212439]	[0.99300325]	0
1	2023-11-17 20:34:17.005	[-1.0603297501, 2.3544967095, -3.5638788326, 5.1387348926, -1.2308457019, -0.7687824608, -3.5881228109, 1.8880837663, -3.2789674274, -3.9563254554, 4.0993439118, -5.6539176395, -0.8775733373, -9.131571192, -0.6093537873, -3.7480276773, -5.0309125017, -0.8748149526, 1.9870535692, 0.7005485718, 0.9204422758, -0.1041491809, 0.3229564351, -0.7418141657, 0.0384120159, 1.0993439146, 1.2603409756, -0.1466244739, -1.4463212439]	[0.99300325]	0
2	2023-11-17 20:34:17.005	[-1.0603297501, 2.3544967095, -3.5638788326, 5.1387348926, -1.2308457019, -0.7687824608, -3.5881228109, 1.8880837663, -3.2789674274, -3.9563254554, 4.0993439118, -5.6539176395, -0.8775733373, -9.131571192, -0.6093537873, -3.7480276773, -5.0309125017, -0.8748149526, 1.9870535692, 0.7005485718, 0.9204422758, -0.1041491809, 0.3229564351, -0.7418141657, 0.0384120159, 1.0993439146, 1.2603409756, -0.1466244739, -1.4463212439]	[0.99300325]	0
3	2023-11-17 20:34:17.005	[-1.0603297501, 2.3544967095, -3.5638788326, 5.1387348926, -1.2308457019, -0.7687824608, -3.5881228109, 1.8880837663, -3.2789674274, -3.9563254554, 4.0993439118, -5.6539176395, -0.8775733373, -9.131571192, -0.6093537873, -3.7480276773, -5.0309125017, -0.8748149526, 1.9870535692, 0.7005485718, 0.9204422758, -0.1041491809, 0.3229564351, -0.7418141657, 0.0384120159, 1.0993439146, 1.2603409756, -0.1466244739, -1.4463212439]	[0.99300325]	0
4	2023-11-17 20:34:17.005	[0.5817662108, 0.097881551, 0.1546819424, 0.4754101949, -0.1978862306, -0.4504344854, 0.0166540447, -0.0256070551, 0.0920561602, -0.2783917153, 0.0593299441, -0.0196585416, -0.4225083157, -0.1217538877, 1.5473094894, 0.2391622864, 0.3553974881, -0.7685165301, -0.7000849355, -0.1190043285, -0.3450517133, -1.1065114108, 0.2523411195, 0.0209441826, 0.2199267436, 0.2540689265, -0.0450225094, 0.1086773898, 0.2547179311]	[0.0010916889]	0

All Inputs Are Tensors

All inputs into an ONNX model must be tensors. This requires that the shape of each element is the same. For example, the following is a proper input:

t [
    [2.35, 5.75],
    [3.72, 8.55],
    [5.55, 97.2]
]

Another example is a 2,2,3 tensor, where the shape of each element is (3,), and each element has 2 rows.

t = [
        [2.35, 5.75, 19.2],
        [3.72, 8.55, 10.5]
    ],
    [
        [5.55, 7.2, 15.7],
        [9.6, 8.2, 2.3]
    ]

In this example each element has a shape of (2,). Tensors with elements of different shapes, known as ragged tensors, are not supported. For example:

t = [
    [2.35, 5.75],
    [3.72, 8.55, 10.5],
    [5.55, 97.2]
])

**INVALID SHAPE**

For models that require ragged tensor or other shapes, see other data formatting options such as Bring Your Own Predict models.

Data Type Consistency

All inputs into an ONNX model must have the same internal data type. For example, the following is valid because all of the data types within each element are float32.

t = [
    [2.35, 5.75],
    [3.72, 8.55],
    [5.55, 97.2]
]

The following is invalid, as it mixes floats and strings in each element:

t = [
    [2.35, "Bob"],
    [3.72, "Nancy"],
    [5.55, "Wani"]
]

The following inputs are valid, as each data type is consistent within the elements.

df = pd.DataFrame({
    "t": [
        [2.35, 5.75, 19.2],
        [5.55, 7.2, 15.7],
    ],
    "s": [
        ["Bob", "Nancy", "Wani"],
        ["Jason", "Rita", "Phoebe"]
    ]
})
df

	t	s
0	[2.35, 5.75, 19.2]	[Bob, Nancy, Wani]
1	[5.55, 7.2, 15.7]	[Jason, Rita, Phoebe]

Parameter	Description
Web Site	https://www.tensorflow.org/
Supported Libraries	tensorflow==2.9.1
Framework	Framework.TENSORFLOW aka tensorflow
Runtime	Native aka tensorflow
Supported File Types	SavedModel format as .zip file

IMPORTANT NOTE

These requirements are <strong>not</strong> for Tensorflow Keras models, only for non-Keras Tensorflow models in the SavedModel format.  For Tensorflow Keras deployment in Wallaroo, see the Tensorflow Keras requirements.

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.

Parameter	Description
Web Site	https://www.python.org/
Supported Libraries	python==3.8
Framework	Framework.PYTHON aka python
Runtime	Native aka python

Python models uploaded to Wallaroo are executed as a native runtime.

Note that Python models - aka “Python steps” - are standalone python scripts that use the python libraries natively supported by the Wallaroo platform. These are used for either simple model deployment (such as ARIMA Statsmodels), or data formatting such as the postprocessing steps. A Wallaroo Python model will be composed of one Python script that matches the Wallaroo requirements.

This is contrasted with Arbitrary Python models, also known as Bring Your Own Predict (BYOP) allow for custom model deployments with supporting scripts and artifacts. These are used with pre-trained models (PyTorch, Tensorflow, etc) along with whatever supporting artifacts they require. Supporting artifacts can include other Python modules, model files, etc. These are zipped with all scripts, artifacts, and a requirements.txt file that indicates what other Python models need to be imported that are outside of the typical Wallaroo platform.

Python Models Requirements

Python models uploaded to Wallaroo are Python scripts that must include the wallaroo_json method as the entry point for the Wallaroo engine to use it as a Pipeline step.

This method receives the results of the previous Pipeline step, and its return value will be used in the next Pipeline step.

If the Python model is the first step in the pipeline, then it will be receiving the inference request data (for example: a preprocessing step). If it is the last step in the pipeline, then it will be the data returned from the inference request.

In the example below, the Python model is used as a post processing step for another ML model. The Python model expects to receive data from a ML Model who’s output is a DataFrame with the column dense_2. It then extracts the values of that column as a list, selects the first element, and returns a DataFrame with that element as the value of the column output.

def wallaroo_json(data: pd.DataFrame):
    print(data)
    return [{"output": [data["dense_2"].to_list()[0][0]]}]

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. In the example above, this results the output field as the out.output field in the Wallaroo inference result.

	time	in.tensor	out.output	check_failures
0	2023-06-20 20:23:28.395	[0.6878518042, 0.1760734021, -0.869514083, 0.3..	[12.886651039123535]	0

• Non-Native Frameworks:
• Hugging Face
• PyTorch
• SKLearn
• Tensorflow keras
• XGBoost
• Arbitrary Python (BYOP)
• MLFlow

input_schema = pa.schema([
    pa.field('inputs', pa.string())
])
output_schema = pa.schema([
    pa.field('output', pa.list_(
        pa.list_(
            pa.float64(),
            list_size=128
        ),
    ))
])

input_schema = pa.schema([
    pa.field('inputs', pa.list_(
        pa.list_(
            pa.list_(
                pa.int64(),
                list_size=3
            ),
            list_size=100
        ),
        list_size=100
    )),
    pa.field('top_k', pa.int64()),
])

output_schema = pa.schema([
    pa.field('score', pa.list_(pa.float64(), list_size=2)),
    pa.field('label', pa.list_(pa.string(), list_size=2)),
])

input_schema = pa.schema([
    pa.field('inputs', 
        pa.list_(
            pa.list_(
                pa.list_(
                    pa.int64(),
                    list_size=3
                ),
                list_size=100
            ),
        list_size=100
    )),
    pa.field('threshold', pa.float64()),
    pa.field('mask_threshold', pa.float64()),
    pa.field('overlap_mask_area_threshold', pa.float64()),
])

output_schema = pa.schema([
    pa.field('score', pa.list_(pa.float64())),
    pa.field('label', pa.list_(pa.string())),
    pa.field('mask', 
        pa.list_(
            pa.list_(
                pa.list_(
                    pa.int64(),
                    list_size=100
                ),
                list_size=100
            ),
    )),
])

input_schema = pa.schema([
    pa.field('inputs', pa.list_( #required
        pa.list_(
            pa.list_(
                pa.int64(),
                list_size=3
            ),
            list_size=100
        ),
        list_size=100
    )),
    # pa.field('max_new_tokens', pa.int64()),  # optional
])

output_schema = pa.schema([
    pa.field('generated_text', pa.list_(pa.string())),
])

input_schema = pa.schema([
    pa.field('inputs', 
        pa.list_(
            pa.list_(
                pa.list_(
                    pa.int64(),
                    list_size=3
                ),
                list_size=100
            ),
        list_size=100
    )),
    pa.field('threshold', pa.float64()),
])

output_schema = pa.schema([
    pa.field('score', pa.list_(pa.float64())),
    pa.field('label', pa.list_(pa.string())),
    pa.field('box', 
        pa.list_( # dynamic output, i.e. dynamic number of boxes per input image, each sublist contains the 4 box coordinates 
            pa.list_(
                    pa.int64(),
                    list_size=4
                ),
            ),
    ),
])

input_schema = pa.schema([
    pa.field('question', pa.string()),
    pa.field('context', pa.string()),
    pa.field('top_k', pa.int64()),
    pa.field('doc_stride', pa.int64()),
    pa.field('max_answer_len', pa.int64()),
    pa.field('max_seq_len', pa.int64()),
    pa.field('max_question_len', pa.int64()),
    pa.field('handle_impossible_answer', pa.bool_()),
    pa.field('align_to_words', pa.bool_()),
])

output_schema = pa.schema([
    pa.field('score', pa.float64()),
    pa.field('start', pa.int64()),
    pa.field('end', pa.int64()),
    pa.field('answer', pa.string()),
])

input_schema = pa.schema([
    pa.field('prompt', pa.string()),
    pa.field('height', pa.int64()),
    pa.field('width', pa.int64()),
    pa.field('num_inference_steps', pa.int64()), # optional
    pa.field('guidance_scale', pa.float64()), # optional
    pa.field('negative_prompt', pa.string()), # optional
    pa.field('num_images_per_prompt', pa.string()), # optional
    pa.field('eta', pa.float64()) # optional
])

output_schema = pa.schema([
    pa.field('images', pa.list_(
        pa.list_(
            pa.list_(
                pa.int64(),
                list_size=3
            ),
            list_size=128
        ),
        list_size=128
    )),
])

input_schema = pa.schema([
    pa.field('inputs', pa.string()),
    pa.field('return_text', pa.bool_()),
    pa.field('return_tensors', pa.bool_()),
    pa.field('clean_up_tokenization_spaces', pa.bool_()),
    # pa.field('extra_field', pa.int64()), # every extra field you specify will be forwarded as a key/value pair
])

output_schema = pa.schema([
    pa.field('summary_text', pa.string()),
])

input_schema = pa.schema([
    pa.field('inputs', pa.string()), # required
    pa.field('top_k', pa.int64()), # optional
    pa.field('function_to_apply', pa.string()), # optional
])

output_schema = pa.schema([
    pa.field('label', pa.list_(pa.string(), list_size=2)), # list with a number of items same as top_k, list_size can be skipped but may lead in worse performance
    pa.field('score', pa.list_(pa.float64(), list_size=2)), # list with a number of items same as top_k, list_size can be skipped but may lead in worse performance
])

input_schema = pa.schema([
    pa.field('inputs', pa.string()), # required
    pa.field('return_tensors', pa.bool_()), # optional
    pa.field('return_text', pa.bool_()), # optional
    pa.field('clean_up_tokenization_spaces', pa.bool_()), # optional
    pa.field('src_lang', pa.string()), # optional
    pa.field('tgt_lang', pa.string()), # optional
    # pa.field('extra_field', pa.int64()), # every extra field you specify will be forwarded as a key/value pair
])

output_schema = pa.schema([
    pa.field('translation_text', pa.string()),
])

input_schema = pa.schema([
    pa.field('inputs', pa.string()), # required
    pa.field('candidate_labels', pa.list_(pa.string(), list_size=2)), # required
    pa.field('hypothesis_template', pa.string()), # optional
    pa.field('multi_label', pa.bool_()), # optional
])

output_schema = pa.schema([
    pa.field('sequence', pa.string()),
    pa.field('scores', pa.list_(pa.float64(), list_size=2)), # same as number of candidate labels, list_size can be skipped by may result in slightly worse performance
    pa.field('labels', pa.list_(pa.string(), list_size=2)), # same as number of candidate labels, list_size can be skipped by may result in slightly worse performance
])

input_schema = pa.schema([
    pa.field('inputs', # required
        pa.list_(
            pa.list_(
                pa.list_(
                    pa.int64(),
                    list_size=3
                ),
                list_size=100
            ),
        list_size=100
    )),
    pa.field('candidate_labels', pa.list_(pa.string(), list_size=2)), # required
    pa.field('hypothesis_template', pa.string()), # optional
]) 

output_schema = pa.schema([
    pa.field('score', pa.list_(pa.float64(), list_size=2)), # same as number of candidate labels
    pa.field('label', pa.list_(pa.string(), list_size=2)), # same as number of candidate labels
])

input_schema = pa.schema([
    pa.field('images', 
        pa.list_(
            pa.list_(
                pa.list_(
                    pa.int64(),
                    list_size=3
                ),
                list_size=640
            ),
        list_size=480
    )),
    pa.field('candidate_labels', pa.list_(pa.string(), list_size=3)),
    pa.field('threshold', pa.float64()),
    # pa.field('top_k', pa.int64()), # we want the model to return exactly the number of predictions, we shouldn't specify this
])

output_schema = pa.schema([
    pa.field('score', pa.list_(pa.float64())), # variable output, depending on detected objects
    pa.field('label', pa.list_(pa.string())), # variable output, depending on detected objects
    pa.field('box', 
        pa.list_( # dynamic output, i.e. dynamic number of boxes per input image, each sublist contains the 4 box coordinates 
            pa.list_(
                    pa.int64(),
                    list_size=4
                ),
            ),
    ),
])

input_schema = pa.schema([
    pa.field('inputs', pa.string()),
    pa.field('return_tensors', pa.bool_()), # optional
    pa.field('return_text', pa.bool_()), # optional
    pa.field('return_full_text', pa.bool_()), # optional
    pa.field('clean_up_tokenization_spaces', pa.bool_()), # optional
    pa.field('prefix', pa.string()), # optional
    pa.field('handle_long_generation', pa.string()), # optional
    # pa.field('extra_field', pa.int64()), # every extra field you specify will be forwarded as a key/value pair
])

output_schema = pa.schema([
    pa.field('generated_text', pa.list_(pa.string(), list_size=1))
])

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

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

pipeline.infer(dataframe)

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

zip -r alohacnnlstm.zip alohacnnlstm/

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

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

pipeline.infer(dataframe)

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

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

List Wallaroo Frameworks

Wallaroo frameworks are listed from the Wallaroo.Framework class. The following demonstrates listing all available supported frameworks.

from wallaroo.framework import Framework

[e.value for e in Framework]

    ['onnx',
    'tensorflow',
    'python',
    'keras',
    'sklearn',
    'pytorch',
    'xgboost',
    'hugging-face-feature-extraction',
    'hugging-face-image-classification',
    'hugging-face-image-segmentation',
    'hugging-face-image-to-text',
    'hugging-face-object-detection',
    'hugging-face-question-answering',
    'hugging-face-stable-diffusion-text-2-img',
    'hugging-face-summarization',
    'hugging-face-text-classification',
    'hugging-face-translation',
    'hugging-face-zero-shot-classification',
    'hugging-face-zero-shot-image-classification',
    'hugging-face-zero-shot-object-detection',
    'hugging-face-sentiment-analysis',
    'hugging-face-text-generation']

Registry Services Roles

Registry service use in Wallaroo typically falls under the following roles.

Model Registry Operations

The following links to guides and information on setting up a model registry (also known as an artifact registry).

Create Model Registry

• Model Registry
• Azure Databricks
• MLflow Model Registry

See Model serving with Azure Databricks for setting up a model registry service using Azure Databricks.

The following steps create an Access Token used to authenticate to an Azure Databricks Model Registry.

1. Log into the Azure Databricks workspace.
2. From the upper right corner access the User Settings.
3. From the Access tokens, select Generate new token.
4. Specify any token description and lifetime. Once complete, select Generate.
5. Copy the token and store in a secure place. Once the Generate New Token module is closed, the token will not be retrievable.

The MLflow Model Registry provides a method of setting up a model registry service. Full details can be found at the MLflow Registry Quick Start Guide.

A generic MLFlow model registry requires no token.

Wallaroo Registry Operations

• Connect Model Registry to Wallaroo: This details the link and connection information to a existing MLFlow registry service. Note that this does not create a MLFlow registry service, but adds the connection and credentials to Wallaroo to allow that MLFlow registry service to be used by other entities in the Wallaroo instance.
• Add a Registry to a Workspace: Add the created Wallaroo Model Registry so make it available to other workspace members.
• Remove a Registry from a Workspace: Remove the link between a Wallaroo Model Registry and a Wallaroo workspace.

Connect Model Registry to Wallaroo

MLFlow Registry connection information is added to a Wallaroo instance through the following endpoint.

• REQUEST URL
  • v1/api/models/create_registry
• PARAMETERS
  • workspace_id (Integer Required): The numerical ID of the workspace to create the registry in.
  • name (String Required): The name for the registry. Registry names are not unique.
  • url (String Required): The full URL of the registry service. For example: https://registry.wallaroo.ai
  • token (String Required): The authentication token used by the registry service.
• RETURNS
  • id (String): The UUID of the registry.
  • workspace_id: The numerical ID of the workspace the registry was connected to.

Connect Model Registry to Wallaroo Example

The following registry will be added to the workspace with the id 1.

import requests
token = "abcdefg"

x = requests.post("https://{APIURL}/v1/api/models/create_registry", 
                    json={
                        "workspace_id": 1, 
                        "name": "sample registry", 
                        "url": "https://registry.wallaroo.ai", 
                        "token": token
                        }, 
                        headers=wl.auth.auth_header()
                    )

{'id': '98f9ca1d-c4e7-4d70-8df4-05c25a64be29', 'workspace_id': 1}

Add Wallaroo Registry Service to Workspace

Registries are assigned to a Wallaroo workspace with the following endpoint. This allows members of the workspace to access the registry connection. A registry can be associated with one or more workspaces.

• REQUEST URL
  • v1/api/models//v1/api/models/attach_registry_to_workspace
• PARAMETERS
  • workspace_id (Integer Required): The numerical ID of the workspace to create the registry in.
  • registry_id (String Required): The ID of the registry in UUID format.
• RETURNS
  • id (String): The UUID of the registry.
  • workspace_id: The numerical ID of the workspace the registry was connected to.

Add Wallaroo Registry Service to Workspace Example

import requests
token = "abcdefg"

x = requests.post("https://{APIURL}/v1/api/models/attach_registry_to_workspace", 
                    json={
                      'id': '98f9ca1d-c4e7-4d70-8df4-05c25a64be29', 
                      'workspace_id': 1},
                        headers=wl.auth.auth_header()
                    )

{
  "registry_id": "3fa85f64-5717-4562-b3fc-2c963f66afa6",
  "workspace_id": 0
}

Remove Wallaroo Registry from a Workspace

Registries are removed from a registry with the following endpoint. This does not remove the registry connection information from the Wallaroo instance, merely removes the association between the registry and that particular workspace.

• REQUEST URL
  • v1/api/models//v1/api/models/remove_registry_from_workspace
• PARAMETERS
  • workspace_id (Integer Required): The numerical ID of the workspace to remove the registry from.
  • registry_id (String Required): The ID of the registry in UUID format.
• RETURNS
  • id (String): The UUID of the registry.
  • workspace_id: The numerical ID of the workspace the registry was connected to.

Remove Wallaroo Registry from a Workspace Example

import requests
token = "abcdefg"

x = requests.post("https://{APIURL}/v1/api/models/remove_registry_from_workspace", 
                    json={
                      'id': '98f9ca1d-c4e7-4d70-8df4-05c25a64be29', 
                      'workspace_id': 1
                      },
                        headers=wl.auth.auth_header()
                    )

{
  "registry_id": "3fa85f64-5717-4562-b3fc-2c963f66afa6",
  "workspace_id": 0
}

Get Registry Details

• REQUEST URL
  • v1/api/models/get_registry
• PARAMETERS
  • id (String Required): The registry ID in UUID format.
• RETURNS
  • id (String): The UUID of the registry.
  • name (String Required): The name for the registry. Registry names are not unique.
  • url (String Required): The full URL of the registry service. For example: https://registry.wallaroo.ai
  • token (String Required): The authentication token used by the registry service.
  • created_at (String Required): The creation date in DateTime format.
  • updated_at (String Required): The updated date in DateTime format.

Get Registry Details Example

The following example demonstrates retrieving the registry with id 98f9ca1d-c4e7-4d70-8df4-05c25a64be29.

import requests
x = requests.post("https://{APIURL}/v1/api/models/get_registry", 
                    json={
                        "registry_id": '98f9ca1d-c4e7-4d70-8df4-05c25a64be29'
                        }, 
                    headers=wl.auth.auth_header()
                )
{
    'id': '98f9ca1d-c4e7-4d70-8df4-05c25a64be29',
    'name': 'sample registry',
    'url': 'https://registry.wallaroo.ai',
    'token': 'dapi67c8c0b04606f730e78b7ae5e3221015-3',
    'created_at': '2023-06-23T15:37:38.38427+00:00',
    'updated_at': '2023-06-23T15:37:38.38427+00:00'
}

Wallaroo Registry Model Operations

List Registries in Workspace

A list of registries associates with a specific workspace are retrieved from the following endpoint.

• REQUEST URL
  • POST /v1/api/models/list_registries
• PARAMETERS
  • workspace-id (String Required): The numerical id of the workspace.
• RETURNS
  • A list of registries with the following fields.
    • created_at (String): The UTC DateTime stamp of when the registry was created.
    • id (String): The id of the Wallaroo registry in UUID format.
    • name (String): The assigned name of the registry.
    • token (String): The authentication token to the model registry.
    • updated_at (String): The UTC DateTime stamp of when the registry was last updated.
    • url (String): The URL to the registry service.

List Registries in Workspace Example

import requests
x = requests.post("{APIURL}v1/api/models/list_registries", 
                    json={
                      "workspace_id": 1
                      }, 
                      headers=wl.auth.auth_header()
                      )
x.json()

[
  {
    "created_at": "2023-07-11T15:21:24.403Z",
    "id": "3fa85f64-5717-4562-b3fc-2c963f66afa6",
    "name": "sample registry",
    "token": "string",
    "updated_at": "2023-07-11T15:21:24.403Z",
    "url": "string"
  }
]

List Models in a Registry

A List of models available to the Wallaroo instance through the MLFlow Registry is performed with the following endpoint.

• REQUEST URL
  • v1/api/models/list_registry_models
• PARAMETERS
  • id (String Required): The registry ID in UUID format.
• RETURNS
  • next_page_token (String): Used to retrieve the next List of models. If None, there are no additional models to list after this request.
  • registered_models (List): List of Models based on the mlflow.entities.model_registry.ModelVersion specification.
    • name (String): The name of the model.
    • user_id (String): The ID of the user.
    • latest_versions (List): The list of other versions of the model. If there are no other versions of the model, this will be None. Each version has the following fields.
      • name (String): The name of the artifact.
      • description (String): The description of the artifact.
      • version (Integer): The version number. Other versions may be removed from the registry, so it is possible for a version to be higher than 1 even if there are no other versions still stored in the registry.
      • status (String): The current status of the model.
      • run_id (String): The run id from the MLFlow service that created the model.
      • run_link (String): Link to the run from the MLFlow service that created the model.
      • source (String): The URL for the specific version artifact on this registry service. For example: 'dbfs:/databricks/mlflow-tracking/123456/abcdefg/artifacts/random_forest_model'.
      • current_stage (String): The current stage of the model version.
      • creation_timestamp (Integer): Creation timestamp in milliseconds since the Unix epoch.
      • last_updated_timestamp (Integer): Last updated timestamp in milliseconds since the Unix epoch.

List Models in a Registry Example

import requests
x = requests.post("{APIURL}v1/api/models/list_registry_models", 
                    json={
                      "registry_id": id
                      }, 
                      headers=wl.auth.auth_header()
                      )
x.json()

{
    'next_page_token': None,
    'registered_models': [
        {
            'name': 'testmodel', 
            'user_id': 'sample.usersj@wallaroo.ai', 
            'latest_versions': None, 
            'creation_timestamp': 1686940722329, 
            'last_updated_timestamp': 1686940722329
        },
        {
            'name': 'testmodel2', 
            'user_id': 'sample.user@wallaroo.ai', 
            'latest_versions': None, 
            'creation_timestamp': 1686940864528, 
            'last_updated_timestamp': 1686940864528
        },
        {
            'name': 'wine_quality', 
            'user_id': 'sample.user@wallaroo.ai', 
            'latest_versions': [
                {
                    'name': 'wine_quality', 
                    'description': None, 
                    'version': '1', 
                    'status': 'READY', 
                    'run_id': 'abcdefg', 
                    'run_link': None, 
                    'source': 'dbfs:/databricks/mlflow-tracking/abcdefg/abcdefg/artifacts/random_forest_model', 
                    'current_stage': 'Archived', 
                    'creation_timestamp': 1686942353367, 
                    'last_updated_timestamp': 1686942597509
                },
                {
                    'name': 'wine_quality', 
                    'description': None, 
                    'version': '2', 
                    'status': 'READY', 
                    'run_id': 'abcdefg', 
                    'run_link': None, 
                    'source': 'dbfs:/databricks/mlflow-tracking/abcdefg/abcdefg/artifacts/model', 
                    'current_stage': 'Production', 
                    'creation_timestamp': 1686942576120, 
                    'last_updated_timestamp': 1686942597646
                }
            ], 
            'creation_timestamp': 1686942353127, 
            'last_updated_timestamp': 1686942597646
        }
    ]
}

List Model Version Artifacts

The artifacts of a specific model version are retrieved through the following endpoint.

• REQUEST URL
  • v1/api/models/list_registry_model_version_artifacts
• PARAMETERS
  • registry_id (String Required): The registry ID in UUID format.
  • name (String Required): The name of the model to retrieve artifacts from.
  • version (String Required): The version of the model to retrieve artifacts from.
• RETURNS
  • List of artifacts with the following fields.
    • file_size (Integer): The size of the file in bytes.
    • full_path (String): The full path to the artifact. For example: https://adb-5939996465837398.18.azuredatabricks.net/api/2.0/dbfs/read?path=/databricks/mlflow-registry/9f38797c1dbf4e7eb229c4011f0f1f18/models/testmodel2/model.pkl
    • is_dir (Boolean): Whether the artifact is a directory or file.
    • modification_time (Integer): Last modification timestamp in milliseconds since the Unix epoch
    • path (String): Relative path to the artifact within the registry. For example: /databricks/mlflow-registry/9f38797c1dbf4e7eb229c4011f0f1f18/models/testmodel2/model.pkl

List Model Version Artifacts Example

import requests
x = requests.post("{APIURL}v1/api/models//list_registry_model_version_artifacts", 
                    json={
                      "name": "wine_quality",
                      "registry_id": "3fa85f64-5717-4562-b3fc-2c963f66afa6",
                      "version": "2"
                      }, 
                      headers=wl.auth.auth_header()
                      )
x.json()

[
  {
    "file_size": 156456,
    "full_path": "https://adb-5939996465837398.18.azuredatabricks.net/api/2.0/dbfs/read?path=/databricks/mlflow-registry/9f38797c1dbf4e7eb229c4011f0f1f18/models/testmodel2/model.pkl",
    "is_dir": false,
    "modification_time": 1686942597646,
    "path": "/databricks/mlflow-registry/9f38797c1dbf4e7eb229c4011f0f1f18/models/testmodel2/model.pkl"
  }
]

Upload Model from Registry

Models are uploaded from a model registry configured in Wallaroo through the following endpoint. The specific artifact that is the model to be deployed is the item to upload to the Wallaroo workspace. Models must comply with Wallaroo model framework and versions as defined in Artifact Requirements.

• REQUEST URL
  • v1/api/models/upload_from_registry
• PARAMETERS
  • registry_id (String Required): The registry ID in UUID format.
  • name (String Required): The name to assign the model in Wallaroo. Model names map onto Kubernetes objects, and must be DNS compliant. The strings for model names must be ASCII alpha-numeric characters or dash (-) only. . and _ are not allowed.
  • path (String Required): URL of the model to upload as specified in List Models in a Registry source field.
  • visibility (String Required): Whether the model is public or private.
  • workspace_id (Integer Required): The numerical id of the workspace to upload the model to.
  • conversion (String Required): The conversion parameters that include the following:
    • framework (String Required): The framework of the model being uploaded. See the list of supported models for more details.
    • python_version (String Required): The version of Python required for model.
    • requirements (String Required): Required libraries. Can be [] if the requirements are default Wallaroo JupyterHub libraries.
    • input_schema (String Optional): The input schema from the Apache Arrow pyarrow.lib.Schema format, encoded with base64.b64encode. Only required for non-native runtime models.
    • output_schema (String Optional): The output schema from the Apache Arrow pyarrow.lib.Schema format, encoded with base64.b64encode. Only required for non-native runtime models.
• RETURNS
  • model_id (String): The numerical id of the model uploaded

Upload Model from Registry Example

import requests
x = requests.post("{APIURL}v1/api/models/upload_from_registry", json={
    "registry_id": id, 
    "name": "uploaded-model-name",
    "path": "<DBFS URL from list_artifacts() call here>",
    "visibility": "public",
    "workspace_id": 1,
    "conversion": {
        "framework": "sklearn",
        "requirements": [],
    },
    "input_schema": "<base64-encoded input schema here>",
    "output_schema": "<base64-encoded output schema here>"
}, headers=wl.auth.auth_header())
x.json()

{'model_id': 34}