diff --git a/backend/python/exllama/Makefile b/backend/python/exllama/Makefile index 52a180ae..f3938b13 100644 --- a/backend/python/exllama/Makefile +++ b/backend/python/exllama/Makefile @@ -8,3 +8,6 @@ run: @echo "Running exllama..." bash run.sh @echo "exllama run." + +test: + bash test.sh \ No newline at end of file diff --git a/backend/python/exllama/test.py b/backend/python/exllama/test.py new file mode 100644 index 00000000..4f909487 --- /dev/null +++ b/backend/python/exllama/test.py @@ -0,0 +1,63 @@ +""" +A test script to test the gRPC service +""" +import unittest +import subprocess +import time +import backend_pb2 +import backend_pb2_grpc + +import grpc + + +class TestBackendServicer(unittest.TestCase): + """ + TestBackendServicer is the class that tests the gRPC service + """ + def setUp(self): + """ + This method sets up the gRPC service by starting the server + """ + self.service = subprocess.Popen(["python3", "exllama.py", "--addr", "localhost:50051"]) + + def tearDown(self) -> None: + """ + This method tears down the gRPC service by terminating the server + """ + self.service.kill() + self.service.wait() + + def test_server_startup(self): + """ + This method tests if the server starts up successfully + """ + time.sleep(10) + try: + self.setUp() + with grpc.insecure_channel("localhost:50051") as channel: + stub = backend_pb2_grpc.BackendStub(channel) + response = stub.Health(backend_pb2.HealthMessage()) + self.assertEqual(response.message, b'OK') + except Exception as err: + print(err) + self.fail("Server failed to start") + finally: + self.tearDown() + + def test_load_model(self): + """ + This method tests if the model is loaded successfully + """ + time.sleep(10) + try: + self.setUp() + with grpc.insecure_channel("localhost:50051") as channel: + stub = backend_pb2_grpc.BackendStub(channel) + response = stub.LoadModel(backend_pb2.ModelOptions(Model="iambestfeed/open_llama_3b_4bit_128g")) + self.assertTrue(response.success) + self.assertEqual(response.message, "Model loaded successfully") + except Exception as err: + print(err) + self.fail("LoadModel service failed") + finally: + self.tearDown() diff --git a/backend/python/exllama/test.sh b/backend/python/exllama/test.sh new file mode 100644 index 00000000..072fb85b --- /dev/null +++ b/backend/python/exllama/test.sh @@ -0,0 +1,14 @@ +#!/bin/bash + +## +## A bash script wrapper that runs the diffusers server with conda + +export PATH=$PATH:/opt/conda/bin + +# Activate conda environment +source activate transformers + +# get the directory where the bash script is located +DIR="$( cd "$( dirname "${BASH_SOURCE[0]}" )" >/dev/null 2>&1 && pwd )" + +python -m unittest $DIR/test.py \ No newline at end of file