def test_create_server(self):
# Create a TritonServerConfig
server_config = TritonServerConfig()
server_config['model-repository'] = MODEL_REPOSITORY_PATH
gpus = ['all']
# Run for both types of environments
self.server = TritonServerFactory.create_server_docker(
image=TRITON_IMAGE, config=server_config, gpus=gpus)
self.server = TritonServerFactory.create_server_local(
path=TRITON_LOCAL_BIN_PATH, config=server_config)
# Try to create a server without specifying model repository and expect
# error
server_config['model-repository'] = None
with self.assertRaises(
AssertionError,
msg="Expected AssertionError for trying to create"
"server without specifying model repository."):
self.server = TritonServerFactory.create_server_docker(
image=TRITON_IMAGE, config=server_config, gpus=gpus)
with self.assertRaises(
AssertionError,
msg="Expected AssertionError for trying to create"
"server without specifying model repository."):
self.server = TritonServerFactory.create_server_local(
path=TRITON_LOCAL_BIN_PATH, config=server_config)
def _test_cpu_stats(self, gpus):
device_requests = [device.device_id() for device in gpus]
# Create a TritonServerConfig
server_config = TritonServerConfig()
server_config['model-repository'] = MODEL_REPOSITORY_PATH
# Test local server cpu_stats
self.server = TritonServerFactory.create_server_local(
path=TRITON_LOCAL_BIN_PATH, config=server_config, gpus=gpus)
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(
f'{TRITON_DOCKER_BIN_PATH} {server_config.to_cli_string()}',
MODEL_REPOSITORY_PATH, TRITON_IMAGE, device_requests, gpus, 8000,
8001, 8002)
_, _ = self.server.cpu_stats()
self.server_docker_mock.assert_cpu_stats_called()
self.server.stop()
def start_stop_docker_args(self):
device_requests, gpu_uuids = self._find_correct_gpu_settings(
self._sys_gpus)
# Create a TritonServerConfig
server_config = TritonServerConfig()
server_config['model-repository'] = MODEL_REPOSITORY_PATH
# Create mounts and labels
mounts = [
'/host/path:/dest/path:ro', '/another/host/path:/some/dest/path:rw'
]
labels = {'RUNNER_ID': 'TEST_RUNNER_ID'}
environment = {'VARIABLE': 'VALUE'}
# Create server in docker, start , wait, and stop
self.server = TritonServerFactory.create_server_docker(
image=TRITON_IMAGE,
config=server_config,
gpus=self._sys_gpus,
mounts=mounts,
labels=labels)
# Start server check that mocked api is called
self.server.start(env=environment)
self.server_docker_mock.assert_server_process_start_called_with(
f"{TRITON_DOCKER_BIN_PATH} {server_config.to_cli_string()}",
MODEL_REPOSITORY_PATH, TRITON_IMAGE, device_requests, gpu_uuids,
8000, 8001, 8002, mounts, labels)
# Stop container and check api calls
self.server.stop()
self.server_docker_mock.assert_server_process_terminate_called()
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_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 setUp(self):
# Mocks
self.mock_server_docker = MockServerDockerMethods()
self.tritonclient_mock = MockTritonClientMethods()
# Create server config
self.server_config = TritonServerConfig()
self.server_config['model-repository'] = MODEL_REPOSITORY_PATH
self.server_config['model-control-mode'] = 'explicit'
# Set CUDA_VISIBLE_DEVICES
os.environ['CUDA_VISIBLE_DEVICES'] = '0'
# Create and start the server
self.server = TritonServerFactory.create_server_docker(
model_path=MODEL_LOCAL_PATH,
image=TRITON_IMAGE,
config=self.server_config)
def _test_start_stop_gpus(self, gpus):
device_requests = [device.device_id() for device in gpus]
# Create a TritonServerConfig
server_config = TritonServerConfig()
server_config['model-repository'] = MODEL_REPOSITORY_PATH
# 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(
f"{TRITON_DOCKER_BIN_PATH} {server_config.to_cli_string()}",
MODEL_REPOSITORY_PATH, TRITON_IMAGE, device_requests, gpus, 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, gpus=gpus)
# 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
],
gpus=gpus)
self.server.stop()
self.server_local_mock.assert_server_process_terminate_called()
def _test_get_logs(self, gpus):
# Create a TritonServerConfig
server_config = TritonServerConfig()
server_config['model-repository'] = MODEL_REPOSITORY_PATH
# Check docker server logs
self.server = TritonServerFactory.create_server_docker(
image=TRITON_IMAGE, config=server_config, gpus=gpus)
self.server.start()
self.server.stop()
self.server_docker_mock.assert_server_process_terminate_called()
self.assertEqual(self.server.logs(), "Triton Server Test Log")
# Create local server logs
self.server = TritonServerFactory.create_server_local(
path=TRITON_LOCAL_BIN_PATH, config=server_config, gpus=gpus)
self.server.start()
self.server.stop()
self.server_local_mock.assert_server_process_terminate_called()
self.assertEqual(self.server.logs(), "Triton Server Test Log")
def setUp(self):
# GPUs
gpus = [
GPUDevice('TEST_DEVICE_NAME', 0, "TEST_PCI_BUS_ID", "TEST_UUID")
]
# Mocks
self.server_docker_mock = MockServerDockerMethods()
self.tritonclient_mock = MockTritonClientMethods()
self.server_docker_mock.start()
self.tritonclient_mock.start()
# Create server config
self.server_config = TritonServerConfig()
self.server_config['model-repository'] = MODEL_REPOSITORY_PATH
self.server_config['model-control-mode'] = 'explicit'
# Set CUDA_VISIBLE_DEVICES
os.environ['CUDA_VISIBLE_DEVICES'] = '0'
# Create and start the server
self.server = TritonServerFactory.create_server_docker(
image=TRITON_IMAGE, config=self.server_config, gpus=gpus)