Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Validate the memory requested for the infer request is not out of bounds #7083

Merged
merged 11 commits into from
Apr 12, 2024
Prev Previous commit
Next Next commit
add test
  • Loading branch information
jbkyang-nvi committed Apr 6, 2024
commit 2f8bf9178c08a4e2ef0d056ff4cdb3b73c8c9978
29 changes: 27 additions & 2 deletions qa/L0_shared_memory/shared_memory_test.py
Original file line number Diff line number Diff line change
Expand Up @@ -163,6 +163,7 @@ def _basic_inference(
big_shm_name="",
big_shm_size=64,
shm_output_offset=0,
shm_output_byte_size=0,
):
input0_data = np.arange(start=0, stop=16, dtype=np.int32)
input1_data = np.ones(shape=16, dtype=np.int32)
Expand Down Expand Up @@ -192,8 +193,12 @@ def _basic_inference(
else:
inputs[1].set_shared_memory("input1_data", 64)

outputs[0].set_shared_memory("output0_data", 64, offset=shm_output_offset)
outputs[1].set_shared_memory("output1_data", 64, offset=shm_output_offset)
outputs[0].set_shared_memory(
"output0_data", 64, offset=shm_output_offset, byte_size=shm_output_byte_size
)
outputs[1].set_shared_memory(
"output1_data", 64, offset=shm_output_offset, byte_size=shm_output_byte_size
)

try:
results = triton_client.infer(
Expand Down Expand Up @@ -348,6 +353,26 @@ def test_infer_offset_out_of_bound(self):
self.assertIn("Invalid offset for shared memory region", error_msg[0])
self._cleanup_server(shm_handles)

def test_infer_byte_size_out_of_bound(self):
Tabrizian marked this conversation as resolved.
Show resolved Hide resolved
# Shared memory byte_size outside output region - Throws error
error_msg = []
shm_handles = self._configure_sever()
offset = 2**64 - 1
byte_size = 2**32

self._basic_inference(
shm_handles[0],
shm_handles[1],
shm_handles[2],
shm_handles[3],
error_msg,
shm_output_offset=offset,
shm_output_byte_size=byte_size,
)
self.assertEqual(len(error_msg), 1)
self.assertIn("Invalid byte size for shared memory region", error_msg[0])
self._cleanup_server(shm_handles)


if __name__ == "__main__":
_protocol = os.environ.get("CLIENT_TYPE", "http")
Expand Down
3 changes: 2 additions & 1 deletion qa/L0_shared_memory/test.sh
Original file line number Diff line number Diff line change
Expand Up @@ -50,7 +50,8 @@ for i in \
test_too_big_shm \
test_mixed_raw_shm \
test_unregisterall \
test_infer_offset_out_of_bound; do
test_infer_offset_out_of_bound \
test_infer_byte_size_out_of_bound; do
for client_type in http grpc; do
SERVER_ARGS="--model-repository=`pwd`/models --log-verbose=1 ${SERVER_ARGS_EXTRA}"
SERVER_LOG="./$i.$client_type.server.log"
Expand Down
Loading