Containerized MLFlow Tutorial

The Containerized MLFlow Tutorial demonstrates how to use Wallaroo with containerized MLFlow models.

Setting Private Registry Configuration in Wallaroo

Configure Via Kots

If Wallaroo was installed via kots, use the following procedure to add the private model registry information.

  1. Launch the Kots Administrative Dashboard through a terminal linked to the Kubernetes cluster. Replace the namespace with the one used in your installation.

    kubectl kots admin-console --namespace wallaroo
  2. Launch the dashboard, by default at http://localhost:8800.

  3. From the admin dashboard, select Config -> Private Model Container Registry.

  4. Enable Provide private container registry credentials for model images.

  5. Provide the following:

    1. Registry URL: The URL of the Containerized Model Container Registry. Typically in the format host:port. In this example, the registry for GitHub is used. NOTE: When setting the URL for the Containerized Model Container Registry, only the actual service address is needed. For example: with the full URL of the model as, the URL would be
    2. email: The email address of the user authenticating to the registry service.
    3. username: The username of the user authentication to the registry service.
    4. password: The password of the user authentication or token to the registry service.
  6. Scroll down and select Save config.

  7. Deploy the new version.

Once complete, the Wallaroo instance will be able to authenticate to the Containerized Model Container Registry and retrieve the images.

Configure via Helm

  1. During either the installation process or updates, set the following in the local-values.yaml file:

    1. privateModelRegistry:

      1. enabled: true
      2. secretName: model-registry-secret
      3. registry: The URL of the private registry.
      4. email: The email address of the user authenticating to the registry service.
      5. username: The username of the user authentication to the registry service.
      6. password: The password of the user authentication to the registry service.

      For example:

      # Other settings - DNS entries, etc.
      # The private registry settings
        enabled: true
        secretName: model-registry-secret
        email: "YOUR EMAIL ADDRESS"
        username: "YOUR USERNAME"
        password: "Your Password here"
  2. Install or update the Wallaroo instance via Helm as per the Wallaroo Helm Install instructions.

Once complete, the Wallaroo instance will be able to authenticate to the registry service and retrieve the images.

This tutorial and the assets can be downloaded as part of the Wallaroo Tutorials repository.

MLFlow Inference with Wallaroo Tutorial

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.3.0 containerized models. For information on how to containerize an MLFlow model, see the MLFlow Documentation.

This tutorial assumes that you have a Wallaroo instance, and have either your own containerized model or use the one from the reference and are running this Notebook from the Wallaroo Jupyter Hub service.

See the Wallaroo Private Containerized Model Container Registry Guide for details on how to configure a Wallaroo instance with a private model registry.

MLFlow Data Formats

When using containerized MLFlow models with Wallaroo, the inputs and outputs must be named. For example, the following output:


Would need to be wrapped with the data values named:

[{"prediction": -12.045839810372835}]

A short sample code for wrapping data may be:

output_df = pd.DataFrame(prediction, columns=["prediction"])
return output_df

MLFlow Models and Wallaroo

MLFlow models are composed of two parts: the model, and the flavors. When submitting a MLFlow model to Wallaroo, both aspects must be part of the ML Model included in the container. For full information about MLFlow model structure, see the MLFlow Documentation.

Wallaroo registers the models from container registries. Organizations will either have to make their containers available in a public or through a private Containerized Model Container Registry service. For examples on setting up a private container registry service, see the Docker Documentation “Deploy a registry server”. For more details on setting up a container registry in a cloud environment, see the related documentation for your preferred cloud provider:

For this example, we will be using the MLFlow containers that was registered in a GitHub container registry service in MLFlow Creation Tutorial Part 03: Container Registration. The address of those containers are:

  • postprocess: . Used for format data after the statsmodel inferences.
  • statsmodel: . The statsmodel generated in MLFlow Creation Tutorial Part 01: Model Creation.


Before uploading and running an inference with a MLFlow model in Wallaroo the following will be required:

  • MLFlow Input Schema: The input schema with the fields and data types for each MLFlow model type uploaded to Wallaroo. In the examples below, the data types are imported using the pyarrow library.
  • An installed Wallaroo instance.
  • The following Python libraries installed:
    • os
    • wallaroo: The Wallaroo SDK. Included with the Wallaroo JupyterHub service by default.

IMPORTANT NOTE: Wallaroo supports MLFlow 1.3.0. Please ensure the MLFlow models used in Wallaroo meet this specification.

MLFlow Inference Steps

To register a containerized MLFlow ML Model into Wallaroo, use the following general step:

  • Import Libraries
  • Connect to Wallaroo
  • Set MLFlow Input Schemas
  • Register MLFlow Model
  • Create Pipeline and Add Model Steps
  • Run Inference

Import Libraries

We start by importing the libraries we will need to connect to Wallaroo and use our MLFlow models. This includes the wallaroo libraries, pyarrow for data types, and the json library for handling JSON data.

import wallaroo
from wallaroo.object import EntityNotFoundError
import pyarrow as pa
import pandas as pd

Connect to the Wallaroo Instance

The first step is to connect to Wallaroo through the Wallaroo client. The Python library is included in the Wallaroo install and available through the Jupyter Hub interface provided with your Wallaroo environment.

This is accomplished using the wallaroo.Client() command, which provides a URL to grant the SDK permission to your specific Wallaroo environment. When displayed, enter the URL into a browser and confirm permissions. Store the connection into a variable that can be referenced later.

If logging into the Wallaroo instance through the internal JupyterHub service, use wl = wallaroo.Client(). For more information on Wallaroo Client settings, see the Client Connection guide.

# Login through local Wallaroo instance

wl = wallaroo.Client()
prefix = 'mlflow'
workspace_name= f"{prefix}statsmodelworkspace"
pipeline_name = f"{prefix}statsmodelpipeline"

mlflowworkspace = wl.get_workspace(name=workspace_name, create_if_not_exist=True)

pipeline = wl.build_pipeline(pipeline_name)

Set MLFlow Input Schemas

Set the MLFlow input schemas through the pyarrow library. In the examples below, the input schemas for both the MLFlow model statsmodels-test and the statsmodels-test-postprocess model.

sm_input_schema = pa.schema([
  pa.field('temp', pa.float32()),
  pa.field('holiday', pa.uint8()),
  pa.field('workingday', pa.uint8()),
  pa.field('windspeed', pa.float32())

pp_input_schema = pa.schema([
    pa.field('predicted_mean', pa.float32())

Register MLFlow Model

Use the register_model_image method to register the Docker container containing the MLFlow models.

statsmodelUrl = ""
postprocessUrl = ""

sm_model = wl.register_model_image(
).configure("mlflow", input_schema=sm_input_schema, output_schema=pp_input_schema)
pp_model = wl.register_model_image(
).configure("mlflow", input_schema=pp_input_schema, output_schema=pp_input_schema)

Create Pipeline and Add Model Steps

With the models registered, we can add the MLFlow models as steps in the pipeline. Once ready, we will deploy the pipeline so it is available for submitting data for running inferences.

created2024-04-16 20:07:55.551305+00:00
last_updated2024-04-16 20:07:55.551305+00:00
from wallaroo.deployment_config import DeploymentConfigBuilder

deployment_config = DeploymentConfigBuilder() \
    .cpus(0.25).memory('1Gi') \
    .sidekick_env(sm_model, {"GUNICORN_CMD_ARGS": "--timeout=180 --workers=1"}) \
    .sidekick_env(pp_model, {"GUNICORN_CMD_ARGS": "--timeout=180 --workers=1"}) \

created2024-04-16 20:07:55.551305+00:00
last_updated2024-04-16 20:07:58.413980+00:00
versions34de522d-f4b6-4e77-8333-fe418efcb7bf, 671530fe-92a2-4221-a7c1-69d295310eee
{'status': 'Running',
 'details': [],
 'engines': [{'ip': '',
   'name': 'engine-57c8cd655b-774pv',
   'status': 'Running',
   'reason': None,
   'details': [],
   'pipeline_statuses': {'pipelines': [{'id': 'mlflowstatsmodelpipeline',
      'status': 'Running'}]},
   'model_statuses': {'models': [{'name': 'mlflowpostprocess',
      'sha': '825ebae48014d297134930028ab0e823bc0d9551334b9a4402c87a714e8156b2',
      'status': 'Running',
      'version': 'f0bca0e7-c9f3-4505-8771-1f743e892dea'},
     {'name': 'mlflowstatmodels',
      'sha': '3afd13d9c5070679e284050cd099e84aa2e5cb7c08a788b21d6cb2397615d018',
      'status': 'Running',
      'version': 'a090af7f-440d-4794-b4f0-2f41a6566cf7'}]}}],
 'engine_lbs': [{'ip': '',
   'name': 'engine-lb-d7cc8fc9c-2dzsc',
   'status': 'Running',
   'reason': None,
   'details': []}],
 'sidekicks': [{'ip': '',
   'name': 'engine-sidekick-mlflowstatmodels-25-56db858b5b-xlmq2',
   'status': 'Running',
   'reason': None,
   'details': [],
   'statuses': '\n'},
  {'ip': '',
   'name': 'engine-sidekick-mlflowpostprocess-26-554d99f8b-cbpwq',
   'status': 'Running',
   'reason': None,
   'details': [],
   'statuses': '\n'}]}

Run Inference

Once the pipeline is running, we can submit our data to the pipeline and return our results. Once finished, we will undeploy the pipeline to return the resources back to the cluster.

# Initial container run may need extra time to finish deploying - adding 90 second timeout to compensate
results = pipeline.infer_from_file('./resources/bike_day_eval_engine.df.json', timeout=90)
created2024-04-16 20:07:55.551305+00:00
last_updated2024-04-16 20:07:58.413980+00:00
versions34de522d-f4b6-4e77-8333-fe418efcb7bf, 671530fe-92a2-4221-a7c1-69d295310eee