From 25a88455934a4218321bd2910b0525ac6c8bdb11 Mon Sep 17 00:00:00 2001 From: Pawel Ziecina Date: Fri, 10 Nov 2023 02:40:03 -0800 Subject: [PATCH] Fix for limiting bytes tensor --- CHANGELOG.md | 3 ++- pytriton/proxy/communication.py | 2 +- tests/unit/test_communication_tensor_store.py | 21 ++++++++++++++----- 3 files changed, 19 insertions(+), 7 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 31c761b..c88f75b 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -19,7 +19,8 @@ limitations under the License. ## 0.4.1 (2023-11-09) - New: Place where workspaces with temporary Triton model repositories and communication file sockets can be configured by `$PYTRITON_HOME` environment variable -- Fix: recover handling `KeyboardInterrupt` in `triton.serve()` +- Fix: Recover handling `KeyboardInterrupt` in `triton.serve()` +- Fix: Remove limit for handling bytes dtype tensors - Build scripts update - Added support for arm64 platform builds diff --git a/pytriton/proxy/communication.py b/pytriton/proxy/communication.py index f48441b..8a017e4 100644 --- a/pytriton/proxy/communication.py +++ b/pytriton/proxy/communication.py @@ -138,7 +138,7 @@ def _deserialize_bytes_tensor(encoded_tensor, dtype, order: Literal["C", "F"] = return np.array(strs, dtype=dtype, order=order) -_MAX_DTYPE_DESCR = 8 +_MAX_DTYPE_DESCR = 16 # up to 16 chars in dtype descr; |S2147483647 (2^31-1) with margin _PARTIAL_HEADER_FORMAT = f"<{_MAX_DTYPE_DESCR}scH" diff --git a/tests/unit/test_communication_tensor_store.py b/tests/unit/test_communication_tensor_store.py index 4b23ae8..146182b 100644 --- a/tests/unit/test_communication_tensor_store.py +++ b/tests/unit/test_communication_tensor_store.py @@ -156,12 +156,23 @@ def test_tensor_store_connection_timeout(tmp_path): ], 2, ), + # 2GB bytes array + ( + [ + np.array(b"a" * (2**31 - 1), dtype=bytes), + ], + 1, + ), ), ) def test_tensor_store_get_put_equal(tensor_store, tensors, n_times): for _ in range(n_times): - tensors_ids = tensor_store.put(tensors) - assert len(tensors) == len(tensors_ids) - for tensor, tensor_id in zip(tensors, tensors_ids): - tensor_retrieved = tensor_store.get(tensor_id) - np.testing.assert_equal(tensor, tensor_retrieved) + try: + tensors_ids = tensor_store.put(tensors) + assert len(tensors) == len(tensors_ids) + for tensor, tensor_id in zip(tensors, tensors_ids): + tensor_retrieved = tensor_store.get(tensor_id) + np.testing.assert_equal(tensor, tensor_retrieved) + finally: + for tensor_id in tensors_ids: + tensor_store.release_block(tensor_id)