def test_put_get_dataset(mock_data, use_cluster):
"""test sending and recieving a dataset with 2D tensors
of every datatype
"""
data = mock_data.create_data((10, 10))
# Create a dataset to put
dataset = Dataset("test-dataset")
for index, tensor in enumerate(data):
key = f"tensor_{str(index)}"
dataset.add_tensor(key, tensor)
client = Client(None, use_cluster)
client.put_dataset(dataset)
rdataset = client.get_dataset("test-dataset")
for index, tensor in enumerate(data):
key = f"tensor_{str(index)}"
rtensor = rdataset.get_tensor(key)
np.testing.assert_array_equal(
rtensor,
tensor,
"Dataset returned from get_dataset not the same as sent dataset",
)
def test_set_script_from_file(use_cluster):
sent_script = read_script_from_file()
c = Client(None, use_cluster)
c.set_script_from_file(
"test-script-file", osp.join(file_path, "./data_processing_script.txt")
)
returned_script = c.get_script("test-script-file")
assert sent_script == returned_script
def test_set_model_from_file(mock_model, use_cluster):
try:
mock_model.create_torch_cnn(filepath="./torch_cnn.pt")
c = Client(None, use_cluster)
c.set_model_from_file("file_cnn", "./torch_cnn.pt", "TORCH", "CPU")
returned_model = c.get_model("file_cnn")
with open("./torch_cnn.pt", "rb") as f:
model = f.read()
assert model == returned_model
finally:
os.remove("torch_cnn.pt")
def test_run_script(use_cluster):
data = np.array([[1, 2, 3, 4, 5]])
c = Client(None, use_cluster)
c.put_tensor("script-test-data", data)
c.set_function("one-to-one", one_to_one)
c.run_script("one-to-one", "one_to_one", ["script-test-data"], ["script-test-out"])
out = c.get_tensor("script-test-out")
assert out == 5
def test_bad_SSDB(use_cluster):
ssdb = os.environ["SSDB"]
del os.environ["SSDB"]
os.environ["SSDB"] = "not-an-address:6379;"
with pytest.raises(RedisConnectionError):
c = Client(None, use_cluster)
os.environ["SSDB"] = ssdb
def test_1D_put_get(mock_data, use_cluster):
"""Test put/get_tensor for 1D numpy arrays"""
client = Client(None, use_cluster)
data = mock_data.create_data(10)
send_get_arrays(client, data)
def test_bad_function_execution(use_cluster):
"""Error raised inside function"""
c = Client(None, use_cluster)
c.set_function("bad-function", bad_function)
data = np.array([1, 2, 3, 4])
c.put_tensor("bad-func-tensor", data)
with pytest.raises(RedisReplyError):
c.run_script("bad-function", "bad_function", ["bad-func-tensor"],
["output"])
def test_wrong_model_name(mock_data, mock_model, use_cluster):
"""User requests to run a model that is not there"""
data = mock_data.create_data(1)
model = mock_model.create_torch_cnn()
c = Client(None, use_cluster)
c.set_model("simple_cnn", model, "TORCH", "CPU")
c.put_tensor("input", data[0])
with pytest.raises(RedisReplyError):
c.run_model("wrong_cnn", ["input"], ["output"])
def test_address(use_cluster):
# get env var to set through client init
ssdb = os.environ["SSDB"]
del os.environ["SSDB"]
# client init should fail if SSDB not set
c = Client(address=ssdb, cluster=use_cluster)
# check if SSDB was set anyway
assert os.environ["SSDB"] == ssdb
def test_torch_inference(mock_model, use_cluster):
# get model and set into database
model = mock_model.create_torch_cnn()
c = Client(None, use_cluster)
c.set_model("torch_cnn", model, "TORCH")
# setup input tensor
data = torch.rand(1, 1, 3, 3).numpy()
c.put_tensor("torch_cnn_input", data)
# run model and get output
c.run_model("torch_cnn",
inputs=["torch_cnn_input"],
outputs=["torch_cnn_output"])
out_data = c.get_tensor("torch_cnn_output")
assert out_data.shape == (1, 1, 1, 1)
def test_missing_script_function(use_cluster):
"""User requests to run a function not in the script"""
c = Client(None, use_cluster)
c.set_function("bad-function", bad_function)
with pytest.raises(RedisReplyError):
c.run_script("bad-function", "not-a-function-in-script",
["bad-func-tensor"], ["output"])
def test_wrong_model_name_from_file(mock_data, mock_model, use_cluster):
"""User requests to run a model that is not there
that was loaded from file."""
try:
data = mock_data.create_data(1)
mock_model.create_torch_cnn(filepath="./torch_cnn.pt")
c = Client(None, use_cluster)
c.set_model_from_file("simple_cnn_from_file", "./torch_cnn.pt",
"TORCH", "CPU")
c.put_tensor("input", data[0])
with pytest.raises(RedisReplyError):
c.run_model("wrong_cnn", ["input"], ["output"])
finally:
os.remove("torch_cnn.pt")
def setup_resnet(model, device, batch_size, address, cluster=True):
"""Set and configure the PyTorch resnet50 model for inference
:param model: path to serialized resnet model
:type model: str
:param device: CPU or GPU
:type device: str
:param batch_size: batch size for the Orchestrator (batches of batches)
:type batch_size: int
:param cluster: true if using a cluster orchestrator
:type cluster: bool
"""
client = Client(address=address, cluster=cluster)
client.set_model_from_file("resnet_model", model, "TORCH", device,
batch_size)
client.set_script_from_file("resnet_script",
"../imagenet/data_processing_script.txt",
device)
import numpy as np
import argparse
import time
from smartredis import Client
parser = argparse.ArgumentParser(
description="SmartRedis ensemble producer process.")
parser.add_argument("--redis-port")
args = parser.parse_args()
time.sleep(2)
c = Client(address="127.0.0.1:" + str(args.redis_port), cluster=False)
data = np.random.rand(1, 1, 3, 3)
c.put_tensor("product", data)
import numpy as np
from smartredis import Client
# Connect a SmartRedis client to Redis database
db_address = "127.0.0.1:6379"
client = Client(address=db_address, cluster=False)
# 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")
# save model into an in-memory buffer then string
buffer = io.BytesIO()
torch.jit.save(module, buffer)
str_model = buffer.getvalue()
return str_model
if __name__ == "__main__":
parser = argparse.ArgumentParser(
description="SmartRedis ensemble producer process.")
parser.add_argument("--exchange", action="store_true")
args = parser.parse_args()
# get model and set into database
model = create_torch_cnn()
c = Client(False)
c.set_model("torch_cnn", model, "TORCH")
keyout = os.getenv("SSKEYOUT")
keyin = os.getenv("SSKEYIN")
assert keyout in ["producer_0", "producer_1"]
if keyout == "producer_0":
c.set_data_source("producer_1" if args.exchange else "producer_0")
data = torch.ones(1, 1, 3, 3).numpy()
data_other = -torch.ones(1, 1, 3, 3).numpy()
elif keyout == "producer_1":
c.set_data_source("producer_0" if args.exchange else "producer_1")
data = -torch.ones(1, 1, 3, 3).numpy()
data_other = torch.ones(1, 1, 3, 3).numpy()
def test_run_script_multi(use_cluster):
data = np.array([[1, 2, 3, 4]])
data_2 = np.array([[5, 6, 7, 8]])
c = Client(None, use_cluster)
c.put_tensor("srpt-multi-out-data-1", data)
c.put_tensor("srpt-multi-out-data-2", data_2)
c.set_function("two-to-one", two_to_one)
c.run_script(
"two-to-one",
"two_to_one",
["srpt-multi-out-data-1", "srpt-multi-out-data-2"],
["srpt-multi-out-output"],
)
out = c.get_tensor("srpt-multi-out-output")
expected = np.array([4, 8])
np.testing.assert_array_equal(
out, expected, "Returned array from script not equal to expected result"
)
def test_bad_script_file(use_cluster):
c = Client(None, use_cluster)
with pytest.raises(FileNotFoundError):
c.set_script_from_file("key", "not-a-file")
def test_bad_callable(use_cluster):
"""user provides none callable function to set_function"""
c = Client(None, use_cluster)
with pytest.raises(TypeError):
c.set_function("key", "not-a-file")
def test_set_get_script(use_cluster):
c = Client(None, use_cluster)
sent_script = read_script_from_file()
c.set_script("test-set-script", sent_script)
script = c.get_script("test-set-script")
assert sent_script == script
def test_set_get_function(use_cluster):
c = Client(None, use_cluster)
c.set_function("test-set-function", one_to_one)
script = c.get_script("test-set-function")
sent_script = inspect.getsource(one_to_one)
assert script == sent_script
def test_bad_device(use_cluster):
c = Client(None, use_cluster)
with pytest.raises(TypeError):
c.set_script("key", "some_script", device="not-a-gpu")
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
db_address = "127.0.0.1:6379"
client = Client(address=db_address, cluster=False)
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
def test_get_non_existant_script(use_cluster):
c = Client(None, use_cluster)
with pytest.raises(RedisReplyError):
script = c.get_script("not-a-script")
import io
import os
import numpy as np
import torch
import torch.nn as nn
from smartredis import Client
if __name__ == "__main__":
parser = argparse.ArgumentParser(
description="SmartRedis ensemble consumer process.")
parser.add_argument("--exchange", action="store_true")
args = parser.parse_args()
# get model and set into database
c = Client(False)
keyin = os.getenv("SSKEYIN")
data_sources = keyin.split(",")
data_sources.sort()
assert data_sources == ["producer_0", "producer_1"]
inputs = {
data_sources[0]: np.ones((1, 1, 3, 3)),
data_sources[1]: -np.ones((1, 1, 3, 3)),
}
for key in data_sources:
c.set_data_source(key)
def test_SSDB_not_set(use_cluster):
ssdb = os.environ["SSDB"]
del os.environ["SSDB"]
with pytest.raises(RedisConnectionError):
c = Client(None, use_cluster)
os.environ["SSDB"] = ssdb
import argparse
import os
from smartredis import Client
parser = argparse.ArgumentParser(description="SmartRedis ensemble consumer process.")
parser.add_argument("--redis-port")
args = parser.parse_args()
# get model and set into database
c = Client(address="127.0.0.1:"+str(args.redis_port), cluster=False)
# Incoming entity prefixes are stored as a comma-separated list
# in the env variable SSKEYIN
keyin = os.getenv("SSKEYIN")
data_sources = keyin.split(",")
data_sources.sort()
for key in data_sources:
c.set_data_source(key)
input_exists = c.poll_tensor("product", 100, 100)
db_tensor = c.get_tensor("product")
print(f"Tensor for {key} is:", db_tensor)
def test_set_model(mock_model, use_cluster):
model = mock_model.create_torch_cnn()
c = Client(None, use_cluster)
c.set_model("simple_cnn", model, "TORCH", "CPU")
returned_model = c.get_model("simple_cnn")
assert model == returned_model
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(address=db_address, cluster=False)
# Place the script in the database
client.set_script_from_file(
"test-script-file", osp.join(file_path, "./data_processing_script.txt")
)
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=False)
# 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",