Python

This section details the SmartRedis Python client to demonstrate its general use within SmartSim applications and RedisAI.

Note

The following Python API examples connect to a database at the address:port 127.0.0.1:6379. When replicating the following example, ensure that you use the address:port of your local Redis instance.

Tensors

The Python client can send and receive tensors from the Redis database, where they are stored as RedisAI data structures. Additionally, Python client API functions involving tensor data are compatible with Numpy arrays and do not require other data types.

import numpy as np
from smartredis import Client

# Connect a SmartRedis client to Redis database
db_address = "127.0.0.1:6379"
client = Client(logger_name="example_put_get_tensor.py")

# Send a 2D tensor to the database
key = "2D_array"
array = np.random.randint(-10, 10, size=(10, 10))
client.put_tensor(key, array)

# Retrieve the tensor
returned_array = client.get_tensor("2D_array")

Datasets

The Python client can store and retrieve tensors and metadata in datasets. For further information about datasets, please refer to the Dataset section of the Data Structures documentation page.

The code below shows how to store and retrieve tensors that belong to a DataSet.

import numpy as np

from smartredis import Client, Dataset

# Create two arrays to store in the DataSet
data_1 = np.random.randint(-10, 10, size=(10,10))
data_2 = np.random.randint(-10, 10, size=(20, 8, 2))

# Create a DataSet object and add the two sample tensors
dataset = Dataset("test-dataset")
dataset.add_tensor("tensor_1", data_1)
dataset.add_tensor("tensor_2", data_2)

# Connect SmartRedis client to Redis database
client = Client(logger_name="example_put_get_dataset.py")

# Place the DataSet into the database
client.put_dataset(dataset)

# Retrieve the DataSet from the database
rdataset = client.get_dataset("test-dataset")

# Retrieve a tensor from inside of the fetched
# DataSet
rdata_1 = rdataset.get_tensor("tensor_1")

Models

The SmartRedis clients allow users to set and use a PyTorch, ONNX, TensorFlow, or TensorFlow Lite model in the database. Models can be sent to the database directly from memory or a file. The code below illustrates how a jit-traced PyTorch model can be used with the Python client library.

import io

import torch
import torch.nn as nn

from smartredis import Client

# Taken from https://pytorch.org/docs/master/generated/torch.jit.trace.html
class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        self.conv = nn.Conv2d(1, 1, 3)

    def forward(self, x):
        return self.conv(x)

net = Net()
example_forward_input = torch.rand(1, 1, 3, 3)

# Trace a module (implicitly traces `forward`) and construct a
# `ScriptModule` with a single `forward` method
module = torch.jit.trace(net, example_forward_input)

# Create a buffer of the traced model
buffer = io.BytesIO()
torch.jit.save(module, buffer)
model = buffer.getvalue()

# Connect a SmartRedis client and set the model in the database
client = Client(logger_name="example_model_torch.py")
client.set_model("torch_cnn", model, "TORCH", "CPU")

# Retrieve the model and verify that the retrieved
# model matches the original model.
returned_model = client.get_model("torch_cnn")
assert model == returned_model

# Setup input tensor
data = torch.rand(1, 1, 3, 3).numpy()
client.put_tensor("torch_cnn_input", data)

# Run model and get output
client.run_model("torch_cnn", inputs=["torch_cnn_input"], outputs=["torch_cnn_output"])
out_data = client.get_tensor("torch_cnn_output")

Users can set models from a file, as shown in the code below.

import os

import torch
import torch.nn as nn

from smartredis import Client


# taken from https://pytorch.org/docs/master/generated/torch.jit.trace.html
class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        self.conv = nn.Conv2d(1, 1, 3)

    def forward(self, x):
        return self.conv(x)

# Connect a SmartRedis client
client = Client(logger_name="example_model_file_torch.py")

try:
    net = Net()
    example_forward_input = torch.rand(1, 1, 3, 3)
    # Trace a module (implicitly traces `forward`) and construct a
    # `ScriptModule` with a single `forward` method
    module = torch.jit.trace(net, example_forward_input)

    # Save the traced model to a file
    torch.jit.save(module, "./torch_cnn.pt")

    # Set the model in the Redis database from the file
    client.set_model_from_file("file_cnn", "./torch_cnn.pt", "TORCH", "CPU")

    # Put a tensor in the database as a test input
    data = torch.rand(1, 1, 3, 3).numpy()
    client.put_tensor("torch_cnn_input", data)

    # Run model and retrieve the output
    client.run_model("file_cnn", inputs=["torch_cnn_input"], outputs=["torch_cnn_output"])
    out_data = client.get_tensor("torch_cnn_output")
finally:
    os.remove("torch_cnn.pt")

Scripts

Scripts are a way to store python-executable code in the database. The Python client can send scripts to the dataset from a file or directly from memory.

The code below illustrates how to avoid storing a function in an intermediate file. With this technique, we can define and send a function to the database on the fly.

import numpy as np
import torch

from smartredis import Client

def two_to_one(data, data_2):
    """Sample torchscript script that returns the
    highest elements in both arguments

    Two inputs to one output
    """
    # return the highest element
    merged = torch.cat((data, data_2))
    return merged.max(1)[0]

# Connect a SmartRedis client to the Redis database
db_address = "127.0.0.1:6379"
client = Client(address=db_address, cluster=True, logger_name="example_script.py")

# Generate some test data to feed to the two_to_one function
data = np.array([[1, 2, 3, 4]])
data_2 = np.array([[5, 6, 7, 8]])

# Put the test data into the Redis database
client.put_tensor("script-data-1", data)
client.put_tensor("script-data-2", data_2)

# Put the function into the Redis database
client.set_function("two-to-one", two_to_one)

# Run the script using the test data
client.run_script(
    "two-to-one",
    "two_to_one",
    ["script-data-1", "script-data-2"],
    ["script-multi-out-output"],
)

# Retrieve the output of the test function
out = client.get_tensor("script-multi-out-output")

The code below shows how to set a script from a file. Running the script set from the file uses the same API calls as in the example shown above.

import os.path as osp
from smartredis import Client

# Construct a string holding the script file location
file_path = osp.dirname(osp.abspath(__file__))
script_path = osp.join(file_path, "./data_processing_script.txt")

# Connect to the Redis database
db_address = "127.0.0.1:6379"
client = Client(logger_name="example_script_file.py")

# Place the script in the database
client.set_script_from_file(
    "test-script-file", osp.join(file_path, "./data_processing_script.txt")
)

This file must be a valid Python script. For the example above, the file data_processing_script.txt looks like this:

def pre_process(inp):
    mean = torch.zeros(1).float().to(inp.device)
    mean[0] = 2.0
    temp = inp.float() * mean
    return temp