def test_cpu_stats(self):
server_config = TritonServerConfig()
server_config['model-repository'] = MODEL_REPOSITORY_PATH
gpus = ['all']
# Test local server cpu_stats
self.server = TritonServerFactory.create_server_local(
path=TRITON_LOCAL_BIN_PATH, config=server_config)
self.server.start()
_, _ = self.server.cpu_stats()
self.server_local_mock.assert_cpu_stats_called()
self.server.stop()
# Test docker server cpu stats
self.server = TritonServerFactory.create_server_docker(
image=TRITON_IMAGE, config=server_config, gpus=gpus)
self.server.start()
# The following needs to be called as it resets exec_run return value
self.server_docker_mock.assert_server_process_start_called_with(
TRITON_DOCKER_BIN_PATH + ' ' + server_config.to_cli_string(),
MODEL_REPOSITORY_PATH, TRITON_IMAGE, 8000, 8001, 8002)
_, _ = self.server.cpu_stats()
self.server_docker_mock.assert_cpu_stats_called()
self.server.stop()
def test_server_config(self):
# Create a TritonServerConfig
server_config = TritonServerConfig()
server_config['model-repository'] = MODEL_REPOSITORY_PATH
# Check config initializations
self.assertIsNone(server_config[CONFIG_TEST_ARG],
msg="Server config had unexpected initial"
f"value for {CONFIG_TEST_ARG}")
# Set value
server_config[CONFIG_TEST_ARG] = True
# Test get again
self.assertTrue(server_config[CONFIG_TEST_ARG],
msg=f"{CONFIG_TEST_ARG} was not set")
# Try to set an unsupported config argument, expect failure
with self.assertRaises(TritonModelAnalyzerException,
msg="Expected exception on trying to set"
"unsupported argument in Triton server"
"config"):
server_config['dummy'] = 1
# Reset test arg
server_config[CONFIG_TEST_ARG] = None
# Finally set a couple of args and then check the cli string
for arg, value in CLI_TO_STRING_TEST_ARGS.items():
server_config[arg] = value
cli_string = server_config.to_cli_string()
for argstring in cli_string.split():
# Parse the created string
arg, value = argstring.split('=')
arg = arg[2:]
# Make sure each parsed arg was in test dict
self.assertIn(arg,
CLI_TO_STRING_TEST_ARGS,
msg=f"CLI string contained unknown argument: {arg}")
# Make sure parsed value is the one from dict, check type too
test_value = CLI_TO_STRING_TEST_ARGS[arg]
self.assertEqual(
test_value,
type(test_value)(value),
msg=f"CLI string contained unknown value: {value}")
def test_triton_server_ssl_options(self):
server_config = TritonServerConfig()
triton_server_flags = {
'grpc-use-ssl': '1',
'grpc-use-ssl-mutual': '1',
'grpc-server-cert': 'a',
'grpc-server-key': 'b',
'grpc-root-cert': 'c',
}
server_config.update_config(triton_server_flags)
expected_cli_str = f"--grpc-use-ssl=1 --grpc-use-ssl-mutual=1 "\
f"--grpc-server-cert=a --grpc-server-key=b --grpc-root-cert=c"
self.assertEqual(server_config.to_cli_string(), expected_cli_str)
def test_start_wait_stop_gpus(self):
# Create a TritonServerConfig
server_config = TritonServerConfig()
server_config['model-repository'] = MODEL_REPOSITORY_PATH
os.environ['CUDA_VISIBLE_DEVICES'] = '0'
# Create server in docker, start , wait, and stop
self.server = TritonServerFactory.create_server_docker(
model_path=MODEL_LOCAL_PATH,
image=TRITON_IMAGE,
config=server_config)
# Set mock status_code to error, and generate exception
self._mock_server_wait_for_ready(assert_raises=True)
# Start server check that mocked api is called
self.server.start()
self.server_docker_mock.assert_server_process_start_called_with(
TRITON_DOCKER_BIN_PATH + ' ' + server_config.to_cli_string(),
MODEL_LOCAL_PATH, MODEL_REPOSITORY_PATH, TRITON_IMAGE, 8000, 8001,
8002)
# Mock status code for connected server then stop
self._mock_server_wait_for_ready(assert_raises=False)
# Stop container and check api calls
self.server.stop()
self.server_docker_mock.assert_server_process_terminate_called()
# Create local server which runs triton as a subprocess
self.server = TritonServerFactory.create_server_local(
path=TRITON_LOCAL_BIN_PATH, config=server_config)
self._mock_server_wait_for_ready(assert_raises=True)
# Check that API functions are called
self.server.start()
self.server_local_mock.assert_server_process_start_called_with(cmd=[
TRITON_LOCAL_BIN_PATH, '--model-repository', MODEL_REPOSITORY_PATH
])
self._mock_server_wait_for_ready(assert_raises=False)
self.server.stop()
self.server_local_mock.assert_server_process_terminate_called()
def test_start_stop_gpus(self):
# Create a TritonServerConfig
server_config = TritonServerConfig()
server_config['model-repository'] = MODEL_REPOSITORY_PATH
gpus = ['all']
# Create server in docker, start , wait, and stop
self.server = TritonServerFactory.create_server_docker(
image=TRITON_IMAGE, config=server_config, gpus=gpus)
# Start server check that mocked api is called
self.server.start()
self.server_docker_mock.assert_server_process_start_called_with(
TRITON_DOCKER_BIN_PATH + ' ' + server_config.to_cli_string(),
MODEL_REPOSITORY_PATH, TRITON_IMAGE, 8000, 8001, 8002)
self.server_docker_mock.raise_exception_on_container_run()
with self.assertRaises(TritonModelAnalyzerException):
self.server.start()
self.server_docker_mock.stop_raise_exception_on_container_run()
# Stop container and check api calls
self.server.stop()
self.server_docker_mock.assert_server_process_terminate_called()
# Create local server which runs triton as a subprocess
self.server = TritonServerFactory.create_server_local(
path=TRITON_LOCAL_BIN_PATH, config=server_config)
# Check that API functions are called
self.server.start()
self.server_local_mock.assert_server_process_start_called_with(cmd=[
TRITON_LOCAL_BIN_PATH, '--model-repository', MODEL_REPOSITORY_PATH
])
self.server.stop()
self.server_local_mock.assert_server_process_terminate_called()