How to Manage and Scale ML Models with Workload Orchestration

How to Manage and Scale ML Models with Workload Orchestration | Wallaroo.AI blog

This blog post was originally published on the Microsoft Community Hub and is the last in a five-part series where we explore best practices for deploying machine learning models in production

Throughout this blog series we have seen how we can easily and quickly get our ML models into production, validate them for desired outcomes, proactively monitor for data drift and take swift proactive action to ensure we have optimal model output. As we scale and deploy more models into this production process across multiple cloud environments, Data Scientists and ML Engineers are burdened with spending too many valuable cycles on the data plumbing and repetitive tasks needed just to get models to run and produce business reports – often using tools not designed for ML workloads.

Data engineers are also spending far too many cycles supporting data scientists as they try to run and analyze ML pipelines instead of building robust upstream data pipelines to ensure business continuity. In attempting to achieve value from their ML efforts, they soon find bottlenecks preventing them from realizing the production demands they need.

In this fifth and final blog post in our MLOps Production series, we will go through model workload orchestration and show how to continue the journey for building scale and ease of management for deploying sustainable and value producing models into production.

ML Workload Orchestration flow works within 3 tiers:

TierDescription
ML Workload OrchestrationUser created custom instructions that provide automated processes that follow the same steps every time without error. Orchestrations contain the instructions to be performed, uploaded as a .ZIP file with the instructions, requirements, and artifacts.
TaskInstructions on when to run an Orchestration as a scheduled Task. Tasks can be Run Once, where it creates a single Task Run, or Run Scheduled, where a Task Run is created on a regular schedule based on the Kubernetes cronjob specifications. If a Task is Run Scheduled, it will create a new Task Run every time the schedule parameters are met until the Task is killed.
Task RunThe execution of a task. These validate business operations and successfully identify any unsuccessful task runs. If the Task is Run Once, then only one Task Run is generated. If the Task is a Run Scheduled task, then a new Task Run will be created each time the schedule parameters are met, with each Task Run having its own results and logs.
Fig 1.

Automate Batch Forecasts with Workload Orchestrations

In this example we have a saved Tensorflow LSTM model trained on the popular Airline Passenger Forecasting dataset. We will walk through:

  • Deploying this model to production using Wallaroo
  • Making connections to external datastores to fetch information for inferencing
  • Automating forecasting jobs to be run on-demand or at scheduled intervals

Our first step is to import the Python libraries needed. 

import numpy as np
import pandas as pd
import tensorflow as tf
import wallaroo
from wallaroo.framework import Framework
from wallaroo.deployment_config import DeploymentConfigBuilder
from google.cloud import bigquery
from google.oauth2 import service_account
import datetime
import time
import utils
import timeseries_plot_utils as tsplot
%matplotlib inline
tsplot.sim_forecast_steps(num_steps=3)

[output]

Fig 2.
Fig 3.
Fig 4.

Model Upload and Deployment

First, we get a connection to our Wallaroo instance and set the workspace we want to deploy to.

You will get the following message and by clicking on the URL and selecting Yes you will be authenticated to the instance. 

Please log into the following URL in a web browser:

https://keycloak.demo.pov.wallaroo.io/auth/realms/master/device?user_code=TRUH-DZII

Login successful!


wl = wallaroo.Client()
workspace = utils.get_workspace(wl, "airline-forecast")
_ = wl.set_current_workspace(workspace)

We specify out Tensorflow SavedModel file and the framework (Tensorflow) and the upload_model function takes care of uploading the model and making it available and ready for deployment in our workspace:

framework=Framework.TENSORFLOW
model = wl.upload_model("lstm-forecast", 
                        "models/tf_lstm_model.zip", 
                        framework=framework)

We set any metadata around our pipeline, determine how much hardware we want each deployed instance to have and how many replicas we want deployed. In this example, we deploy 1 but I can deploy more to handle concurrent requests as necessary. Following that, we’re ready to build our pipeline and deploy to production.

pipeline_name = "lstm-forecast-airline"
deployment_config = DeploymentConfigBuilder() \
    .replica_count(1) \
    .cpus(1).memory('2Gi') \
    .build()
pipeline = wl.build_pipeline(pipeline_name) \
    .add_model_step(model) \
    .deploy(deployment_config=deployment_config)

[output] ok

Now, we can test our deployment with an input dataset. We use a quick function to grab one datapoint and scale it:

sample_data_point = utils.get_sample_data()
sample_data_point

[output]

array([[[0.48455598],
        [0.61389961],
        [0.6969112 ]]])

We can then put this into the DataFrame format Wallaroo expects

input_df = pd.DataFrame({"lstm_input": sample_data_point.reshape(1,1,3).tolist()})
input_df

[output]

results = pipeline.infer(input_df)
results

[output]

Connect to Production Data Storage

Now that we can see my model works as expected in production, let’s connect to our production BigQuery instance to source data from there:

conn = wl.get_connection("bigquerybatchstorage")
bq_client = bigquery.Client(
    credentials = service_account.Credentials.from_service_account_info(
        conn.details()), 
    project=conn.details()['project_id']
)

Now that we have this connection, I can write queries to fetch batch data live:

sample_dataframe = bq_client.query(
        f"""
        SELECT *
        FROM sample_datasets.airline_passengers
        LIMIT 5"""
    ).to_dataframe()

[output]

Automate Batch Forecasts with Workload Orchestrations

Now that I have a production model and a connection to my production datasource, all that remains is to encapsulate this process so that we can kickoff batch jobs on a schedule or on demand. This is done using Wallaroo’s Machine Learning Workload Orchestration feature.

The first step here is to create a Python file that encapsulates the workflow. We have done that in our automation/main.py file. Once we have this file and all dependencies in a zip file, I can register this orchestration in Wallaroo:

orchestration = wl.upload_orchestration(name="airline_forecast", path="./automation/forecast_orchestration.zip")

We can confirm we see this orchestration along with any others:

wl.list_orchestrations()[-1]

[output]

We can now create tasks of this orchestration. Tasks are instances of this orchestration that run with a particular set of inputs. These can be created on demand or on schedule:

orchestration = wl.list_orchestrations()[-1]
start_time = datetime.datetime.now()
task = orchestration.run_once(name="forecast_run", json_args={"workspace_name": "airline-forecast","pipeline_name": "lstm-forecast-airline","current_date": "1950-06","forecast_window": 1})

We  can examine the status of the runs manually:

task.last_runs()

[output]

Or we can block by examining log outputs for the pipeline:

logs = pipeline.logs(start_datetime = start_time, end_datetime = datetime.datetime.now())
while logs.empty:
    time.sleep(1)
    logs = pipeline.logs(start_datetime = start_time, end_datetime = datetime.datetime.now())

Finally, we can schedule this job to run on a regular basis using cron-style scheduling. Let’s have this job run every Sunday:

schedule = "0 23 * * 0"
task_scheduled = orchestration.run_scheduled(name="airline-forecast-weekly-scheduled", timeout=600, schedule=schedule, json_args={"workspace_name": "airline-forecast", "pipeline_name": "lstm-forecast-airline","current_date": "1950-06","forecast_window": 1})

Final Product

I can use on-demand runs of this orchestration to create the plots we saw earlier:

tsplot.sim_forecast_steps(num_steps=3)

Conclusion

In this final blog post in our series we have addressed a very common set of challenges that AI teams face with production ML workloads and how to solve them through Model Workload Orchestration. This means that we can easily define, automate, and scale recurring production ML workloads that ingest data from predefined data sources, run inferencing, and deposit the results to a predefined location efficiently and easily. 

If you want to try the steps in this blog post series you can access the tutorial ML Workload Orchestration and use the free inference servers available on the Azure Marketplace. Or you can download a free Wallaroo.AI Community Edition which you can install to GitHub Codespaces, or Azure. 

Wallaroo.AI is a unified production ML platform built for Data Scientists and ML Engineers for easily deploying, observing, and optimizing machine learning in production at scale – in any cloud, on-prem, or at the edge.

 

Author

  • Martin Bald

    Sr Manager Technical Product Marketing & DevRel Community

Table of Contents

Topics

Share

Related Blog Posts

Keep up to date with the latest ML production news sign up for the Wallaroo.AI newsletter

© 2026 Wallaroo.AI. All rights reserved.
Platform
Platform Learn how our unified platform enables ML deployment, serving, observability and optimization
Technology Get a deeper dive into the unique technology behind our ML production platform
Solutions
Solutions See how our unified ML platform supports any model for any use case
Computer Vision (AI) Run even complex models in constrained environments, with hundreds or thousands of endpoints
Resources