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.

 1
 2import numpy as np
 3from smartredis import Client
 4
 5# Connect a SmartRedis client to Redis database
 6db_address = "127.0.0.1:6379"
 7client = Client(logger_name="example_put_get_tensor.py")
 8
 9# Send a 2D tensor to the database
10key = "2D_array"
11array = np.random.randint(-10, 10, size=(10, 10))
12client.put_tensor(key, array)
13
14# Retrieve the tensor
15returned_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.

 1import numpy as np
 2
 3from smartredis import Client, Dataset
 4
 5# Create two arrays to store in the DataSet
 6data_1 = np.random.randint(-10, 10, size=(10,10))
 7data_2 = np.random.randint(-10, 10, size=(20, 8, 2))
 8
 9# Create a DataSet object and add the two sample tensors
10dataset = Dataset("test-dataset")
11dataset.add_tensor("tensor_1", data_1)
12dataset.add_tensor("tensor_2", data_2)
13
14# Connect SmartRedis client to Redis database
15client = Client(logger_name="example_put_get_dataset.py")
16
17# Place the DataSet into the database
18client.put_dataset(dataset)
19
20# Retrieve the DataSet from the database
21rdataset = client.get_dataset("test-dataset")
22
23# Retrieve a tensor from inside of the fetched
24# DataSet
25rdata_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.

 1import io
 2
 3import torch
 4import torch.nn as nn
 5
 6from smartredis import Client
 7
 8# Taken from https://pytorch.org/docs/master/generated/torch.jit.trace.html
 9class Net(nn.Module):
10    def __init__(self):
11        super(Net, self).__init__()
12        self.conv = nn.Conv2d(1, 1, 3)
13
14    def forward(self, x):
15        return self.conv(x)
16
17net = Net()
18example_forward_input = torch.rand(1, 1, 3, 3)
19
20# Trace a module (implicitly traces `forward`) and construct a
21# `ScriptModule` with a single `forward` method
22module = torch.jit.trace(net, example_forward_input)
23
24# Create a buffer of the traced model
25buffer = io.BytesIO()
26torch.jit.save(module, buffer)
27model = buffer.getvalue()
28
29# Connect a SmartRedis client and set the model in the database
30client = Client(logger_name="example_model_torch.py")
31client.set_model("torch_cnn", model, "TORCH", "CPU")
32
33# Retrieve the model and verify that the retrieved
34# model matches the original model.
35returned_model = client.get_model("torch_cnn")
36assert model == returned_model
37
38# Setup input tensor
39data = torch.rand(1, 1, 3, 3).numpy()
40client.put_tensor("torch_cnn_input", data)
41
42# Run model and get output
43client.run_model("torch_cnn", inputs=["torch_cnn_input"], outputs=["torch_cnn_output"])
44out_data = client.get_tensor("torch_cnn_output")

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

 1import os
 2
 3import torch
 4import torch.nn as nn
 5
 6from smartredis import Client
 7
 8
 9# taken from https://pytorch.org/docs/master/generated/torch.jit.trace.html
10class Net(nn.Module):
11    def __init__(self):
12        super(Net, self).__init__()
13        self.conv = nn.Conv2d(1, 1, 3)
14
15    def forward(self, x):
16        return self.conv(x)
17
18# Connect a SmartRedis client
19client = Client(logger_name="example_model_file_torch.py")
20
21try:
22    net = Net()
23    example_forward_input = torch.rand(1, 1, 3, 3)
24    # Trace a module (implicitly traces `forward`) and construct a
25    # `ScriptModule` with a single `forward` method
26    module = torch.jit.trace(net, example_forward_input)
27
28    # Save the traced model to a file
29    torch.jit.save(module, "./torch_cnn.pt")
30
31    # Set the model in the Redis database from the file
32    client.set_model_from_file("file_cnn", "./torch_cnn.pt", "TORCH", "CPU")
33
34    # Put a tensor in the database as a test input
35    data = torch.rand(1, 1, 3, 3).numpy()
36    client.put_tensor("torch_cnn_input", data)
37
38    # Run model and retrieve the output
39    client.run_model("file_cnn", inputs=["torch_cnn_input"], outputs=["torch_cnn_output"])
40    out_data = client.get_tensor("torch_cnn_output")
41finally:
42    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.

 1
 2import numpy as np
 3import torch
 4
 5from smartredis import Client
 6
 7def two_to_one(data, data_2):
 8    """Sample torchscript script that returns the
 9    highest elements in both arguments
10
11    Two inputs to one output
12    """
13    # return the highest element
14    merged = torch.cat((data, data_2))
15    return merged.max(1)[0]
16
17# Connect a SmartRedis client to the Redis database
18db_address = "127.0.0.1:6379"
19client = Client(address=db_address, cluster=True, logger_name="example_script.py")
20
21# Generate some test data to feed to the two_to_one function
22data = np.array([[1, 2, 3, 4]])
23data_2 = np.array([[5, 6, 7, 8]])
24
25# Put the test data into the Redis database
26client.put_tensor("script-data-1", data)
27client.put_tensor("script-data-2", data_2)
28
29# Put the function into the Redis database
30client.set_function("two-to-one", two_to_one)
31
32# Run the script using the test data
33client.run_script(
34    "two-to-one",
35    "two_to_one",
36    ["script-data-1", "script-data-2"],
37    ["script-multi-out-output"],
38)
39
40# Retrieve the output of the test function
41out = 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.

 1
 2import os.path as osp
 3from smartredis import Client
 4
 5# Construct a string holding the script file location
 6file_path = osp.dirname(osp.abspath(__file__))
 7script_path = osp.join(file_path, "./data_processing_script.txt")
 8
 9# Connect to the Redis database
10db_address = "127.0.0.1:6379"
11client = Client(logger_name="example_script_file.py")
12
13# Place the script in the database
14client.set_script_from_file(
15    "test-script-file", osp.join(file_path, "./data_processing_script.txt")
16)

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

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