Workshop Notebook 2: Automation with Wallaroo Connections

Wallaroo Connections are definitions set by MLOps engineers that are used by other Wallaroo users for connection information to a data source.

This provides MLOps engineers a method of creating and updating connection information for data stores: databases, Kafka topics, etc. Wallaroo Connections are composed of three main parts:

• Name: The unique name of the connection.
• Type: A user defined string that designates the type of connection. This is used to organize connections.
• Details: Details are a JSON object containing the information needed to make the connection. This can include data sources, authentication tokens, etc.

Wallaroo Connections are only used to store the connection information used by other processes to create and use external connections. The user still has to provide the libraries and other elements to actually make and use the conneciton.

The primary advantage is Wallaroo connections allow scripts and other code to retrieve the connection details directly from their Wallaroo instance, then refer to those connection details. They don’t need to know what those details actually - they can refer to them in their code to make their code more flexible.

For this step, we will use a Google BigQuery dataset to retrieve the inference information, predict the next month of sales, then store those predictions into another table. This will use the Wallaroo Connection feature to create a Connection, assign it to our workspace, then perform our inferences by using the Connection details to connect to the BigQuery dataset and tables.

Prerequisites

• A Wallaroo instance version 2023.2.1 or greater.

References

• Wallaroo SDK Essentials Guide: Model Uploads and Registrations: Python Models
• Wallaroo SDK Essentials Guide: Pipeline Management
• Wallaroo SDK Essentials Guide: Data Connections Management

Preliminaries

In the blocks below we will preload some required libraries.

For convenience, the following helper functions are defined to retrieve previously created workspaces, models, and pipelines:

• get_workspace(name, client): This takes in the name and the Wallaroo client being used in this session, and returns the workspace matching name. If no workspaces are found matching the name, raises a KeyError and returns None.
• get_model_version(model_name, workspace): Retrieves the most recent model version from the model matching the model_name within the provided workspace. If no model matches that name, raises a KeyError and returns None.
• get_pipeline(pipeline_name, workspace): Retrieves the most pipeline from the workspace matching the pipeline_name within the provided workspace. If no model matches that name, raises a KeyError and returns None.

import json
import os
import datetime

import wallaroo
from wallaroo.object import EntityNotFoundError
from wallaroo.framework import Framework

# used to display dataframe information without truncating
from IPython.display import display
import pandas as pd
import numpy as np

pd.set_option('display.max_colwidth', None)

import time
import pyarrow as pa

## convenience functions from the previous notebooks

# return the workspace called <name> through the Wallaroo client.
def get_workspace(name, client):
    workspace = None
    for ws in client.list_workspaces():
        if ws.name() == name:
            workspace= ws
            return workspace
    # if no workspaces were found
    if workspace==None:
        raise KeyError(f"Workspace {name} was not found.")
    return workspace

# returns the most recent model version in a workspace for the matching `model_name`
def get_model_version(model_name, workspace):
    modellist = workspace.models()
    model_version = [m.versions()[-1] for m in modellist if m.name() == model_name]
    # if no models match, return None
    if len(modellist) <= 0:
        raise KeyError(f"Model {mname} not found in this workspace")
        return None
    return model_version[0]

# get a pipeline by name in the workspace
def get_pipeline(pipeline_name, workspace):
    plist = workspace.pipelines()
    pipeline = [p for p in plist if p.name() == pipeline_name]
    if len(pipeline) <= 0:
        raise KeyError(f"Pipeline {pipeline_name} not found in this workspace")
        return None
    return pipeline[0]

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.

## blank space to log in 

wl = wallaroo.Client()

Set Configurations

Set the workspace, pipeline, and model used from Notebook 1. The helper functions will make this task easier.

Set Configurations References

• Wallaroo SDK Essentials Guide: Workspace Management
• Wallaroo SDK Essentials Guide: Pipeline Management

# retrieve the previous workspace, model, and pipeline version

workspace_name = "workshop-workspace-summarization"

workspace = get_workspace(workspace_name, wl)

# set your current workspace to the workspace that you just created
wl.set_current_workspace(workspace)

model_name = 'hf-summarizer'

prime_model_version = get_model_version(model_name, workspace)

pipeline_name = 'hf-summarizer'

pipeline = get_pipeline(pipeline_name, workspace)

# display workspace, model, and pipeline

display(wl.get_current_workspace())
display(prime_model_version)
display(pipeline)

{'name': 'workshop-workspace-summarization', 'id': 13, 'archived': False, 'created_by': 'b030ff9c-41eb-49b4-afdf-2ccbecb6be5d', 'created_at': '2023-10-05T16:17:44.447738+00:00', 'models': [{'name': 'hf-summarizer', 'versions': 1, 'owner_id': '""', 'last_update_time': datetime.datetime(2023, 10, 5, 16, 19, 43, 718500, tzinfo=tzutc()), 'created_at': datetime.datetime(2023, 10, 5, 16, 19, 43, 718500, tzinfo=tzutc())}], 'pipelines': [{'name': 'hf-summarizer', 'create_time': datetime.datetime(2023, 10, 5, 16, 31, 44, 8667, tzinfo=tzutc()), 'definition': '[]'}]}

Deploy the Pipeline with the Model Version Step

As per the other workshops:

1. Clear the pipeline of all steps.
2. Add the model version as a pipeline step.
3. Deploy the pipeline with the following deployment configuration:

deploy_config = wallaroo.DeploymentConfigBuilder().replica_count(1).cpus(0.5).memory("1Gi").build()

pipeline.clear()
pipeline.add_model_step(prime_model_version)

deployment_config = wallaroo.DeploymentConfigBuilder() \
    .cpus(0.25).memory('1Gi') \
    .sidekick_cpus(prime_model_version, 1) \
    .sidekick_memory(prime_model_version, "4Gi") \
    .build()
pipeline.deploy(deployment_config=deployment_config)

Create the Connection

For this demonstration, the connection set to a specific file on a GitHub repository. The connection details can be anything that can be stored in JSON: connection URLs, tokens, etc.

This connection will set a URL to pull a file from GitHub, then use the file contents to perform an inference.

Wallaroo connections are created through the Wallaroo Client create_connection(name, type, details) method. See the Wallaroo SDK Essentials Guide: Data Connections Management guide for full details.

Note that connection names must be unique across the Wallaroo instance - if needed, use random characters at the end to make sure your connection doesn’t have the same name as a previously created connection.

Here’s an example connection used to retrieve the same CSV file used in ./data/test_data.df.json: https://raw.githubusercontent.com/WallarooLabs/Workshops/main/Linear%20Regression/Real%20Estate/data/test_data.df.json

Create the Connection Exercise

# set the connection information for other steps
# suffix is used to create a unique data connection

# set the connection information for other steps
# suffix is used to create a unique data connection
suffix=''
summary_connection_input_name = f'summary-sample-connection{suffix}'
summary_connection_input_type = "HTTP"
summary_connection_input_argument = { 
    "url": "https://raw.githubusercontent.com/WallarooLabs/Workshops/20230927-foundations-v2/LLM/Summarization/data/test_summarization.df.json"
    }

wl.create_connection(summary_connection_input_name, summary_connection_input_type, summary_connection_input_argument)

# set the connection information for other steps
# suffix is used to create a unique data connection
suffix=''
summary_connection_input_name = f'summary-sample-connection{suffix}'
summary_connection_input_type = "HTTP"
summary_connection_input_argument = { 
    "url": "https://raw.githubusercontent.com/WallarooLabs/Workshops/20230927-foundations-v2/LLM/Summarization/data/test_summarization.df.json"
    }

wl.create_connection(summary_connection_input_name, summary_connection_input_type, summary_connection_input_argument)

List Connections

Connections for the entire Wallaroo instance are listed with Wallaroo Client list_connections() method.

List Connections Exercise

Here’s an example of listing the connections when the Wallaroo client is wl.

wl.list_connections()

# list the connections here

wl.list_connections()

Get Connection by Name

To retrieve a previosly created conneciton, we can assign it to a variable with the method Wallaroo Client.get_connection(connection_name). Then we can display the connection itself. Notice that when displaying a connection, the details section will be hidden, but they are retrieved with connection.details(). Here’s an example:

myconnection = client.get_connection("My amazing connection")
display(myconnection)
display(myconnection.details()

Use that code to retrieve your new connection.

Get Connection by Name Example

Here’s an example based on the Wallaroo client saved as wl.

wl.get_connection(forecast_connection_input_name)

# get the connection by name

this_connection = wl.get_connection(summary_connection_input_name)
this_connection

Add Connection to Workspace

We’ll now add the connection to our workspace so it can be retrieved by other workspace users. The method Workspace add_connection(connection_name) adds a Data Connection to a workspace. The method Workspace list_connections() displays a list of connections attached to the workspace.

Add Connection to Workspace Exercise

Use the connection we just created, and add it to the sample workspace. Here’s a code example where the workspace is saved to the variable workspace and the connection is saved as forecast_connection_input_name.

workspace.add_connection(forecast_connection_input_name)

workspace.add_connection(summary_connection_input_name)
workspace.list_connections()

Retrieve Connection from Workspace

To simulate a data scientist’s procedural flow, we’ll now retrieve the connection from the workspace. Specific connections are retrieved by specifying their position in the returned list.

For example, if we have two connections in a workspace and we want the second one, we can assign it to a variable with list_connections[1].

Create a new variable and retrieve the connection we just assigned to the workspace.

Retrieve Connection from Workspace Exercise

Retrieve the connection that was just associated with the workspace. You’ll use the list_connections method, then assign a variable to the connection. Here’s an example if the connection is the most recently one added to the workspace workspace.

forecast_connection = workspace.list_connections()[-1]

forecast_connection = workspace.list_connections()[-1]
display(forecast_connection)

Run Inference with Connection

Connections can be used for different purposes: uploading new models, engine configurations - any place that data is needed. This exercise will use the data connection to perform an inference through our deployed pipeline.

Run Inference with Connection Exercise

We’ll now retrieve sample data through the Wallaroo connection, and perform a sample inference. The connection details are retrieved through the Connection details() method. Use them to retrieve the pandas record file and convert it to a DataFrame, and use it with our sample model.

Here’s a code example that uses the Python requests library to retrieve the file information, then turns it into a DataFrame for the inference request.

display(forecast_connection.details()['url'])

import requests

response = requests.get(
                    forecast_connection.details()['url']
                )

# display(response.json())

df = pd.DataFrame(response.json())

pipeline.infer(df)

display(forecast_connection.details()['url'])

import requests

response = requests.get(
                    forecast_connection.details()['url']
                )

# display(response.json())

df = pd.DataFrame(response.json())
display(df)

multiple_result = pipeline.infer(df, timeout=60)
display(multiple_result)

'https://raw.githubusercontent.com/WallarooLabs/Workshops/20230927-foundations-v2/LLM/Summarization/data/test_summarization.df.json'

Cleaning up.

Now that the workshop is complete, don’t forget to undeploy your pipeline to free up the resources.

pipeline.undeploy()

Congratulations!

In this workshop you have:

• Deployed a single step house price prediction pipeline and sent data to it.
• Create a new Wallaroo connection.
• Assigned the connection to a workspace.
• Retrieved the connection from the workspace.
• Used the data connection to retrieve information from outside of Wallaroo, and use it for an inference.

Great job!