Python
Contents
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(address=db_address, cluster=True)
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
.
1
2import numpy as np
3
4from smartredis import Client, Dataset
5
6# Create two arrays to store in the DataSet
7data_1 = np.random.randint(-10, 10, size=(10,10))
8data_2 = np.random.randint(-10, 10, size=(20, 8, 2))
9
10# Create a DataSet object and add the two sample tensors
11dataset = Dataset("test-dataset")
12dataset.add_tensor("tensor_1", data_1)
13dataset.add_tensor("tensor_2", data_2)
14
15# Connect SmartRedis client to Redis database
16db_address = "127.0.0.1:6379"
17client = Client(address=db_address, cluster=True)
18
19# Place the DataSet into the database
20client.put_dataset(dataset)
21
22# Retrieve the DataSet from the database
23rdataset = client.get_dataset("test-dataset")
24
25# Retrieve a tensor from inside of the fetched
26# DataSet
27rdata_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.
1
2import io
3
4import torch
5import torch.nn as nn
6
7from smartredis import Client
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
18net = Net()
19example_forward_input = torch.rand(1, 1, 3, 3)
20
21# Trace a module (implicitly traces `forward`) and construct a
22# `ScriptModule` with a single `forward` method
23module = torch.jit.trace(net, example_forward_input)
24
25# Create a buffer of the traced model
26buffer = io.BytesIO()
27torch.jit.save(module, buffer)
28model = buffer.getvalue()
29
30# Connect a SmartRedis client and set the model in the database
31db_address = "127.0.0.1:6379"
32client = Client(address=db_address, cluster=True)
33client.set_model("torch_cnn", model, "TORCH", "CPU")
34
35# Retrieve the model and verify that the retrieved
36# model matches the original model.
37returned_model = client.get_model("torch_cnn")
38assert model == returned_model
39
40# Setup input tensor
41data = torch.rand(1, 1, 3, 3).numpy()
42client.put_tensor("torch_cnn_input", data)
43
44# Run model and get output
45client.run_model("torch_cnn", inputs=["torch_cnn_input"], outputs=["torch_cnn_output"])
46out_data = client.get_tensor("torch_cnn_output")
Users can set models from a file, as shown in the code below.
1
2import os
3
4import torch
5import torch.nn as nn
6
7from smartredis import Client
8
9
10# taken from https://pytorch.org/docs/master/generated/torch.jit.trace.html
11class Net(nn.Module):
12 def __init__(self):
13 super(Net, self).__init__()
14 self.conv = nn.Conv2d(1, 1, 3)
15
16 def forward(self, x):
17 return self.conv(x)
18
19# Connect a SmartRedis client
20db_address = "127.0.0.1:6379"
21client = Client(address=db_address, cluster=True)
22
23try:
24 net = Net()
25 example_forward_input = torch.rand(1, 1, 3, 3)
26 # Trace a module (implicitly traces `forward`) and construct a
27 # `ScriptModule` with a single `forward` method
28 module = torch.jit.trace(net, example_forward_input)
29
30 # Save the traced model to a file
31 torch.jit.save(module, "./torch_cnn.pt")
32
33 # Set the model in the Redis database from the file
34 client.set_model_from_file("file_cnn", "./torch_cnn.pt", "TORCH", "CPU")
35
36 # Put a tensor in the database as a test input
37 data = torch.rand(1, 1, 3, 3).numpy()
38 client.put_tensor("torch_cnn_input", data)
39
40 # Run model and retrieve the output
41 client.run_model("file_cnn", inputs=["torch_cnn_input"], outputs=["torch_cnn_output"])
42 out_data = client.get_tensor("torch_cnn_output")
43finally:
44 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)
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(address=db_address, cluster=True)
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