Fix off by 1 on max offset

triton-inference-server · Feb 26, 2024 · e544dff · e544dff
1 parent e921c12
commit e544dff
Show file tree

Hide file tree

Showing 2 changed files with 10 additions and 4 deletions.
diff --git a/qa/L0_shared_memory/shared_memory_test.py b/qa/L0_shared_memory/shared_memory_test.py
@@ -329,9 +329,12 @@ def test_infer_offset_out_of_bound(self):
         error_msg = []
         shm_handles = self._configure_sever()
         if _protocol == "http":
-            offset = 18446744073709551584  # -0x20 in int64_t
+            # -0x20 in int64_t
+            offset = 18446744073709551584
         else:
-            offset = 9223372036854775807  # grpc will throw if > int64
+            # grpc will throw if > int64, instead test for exceed shm region
+            # size by 1 byte, given its size is 64 bytes
+            offset = 64
         self._basic_inference(
             shm_handles[0],
             shm_handles[1],

diff --git a/src/shared_memory_manager.cc b/src/shared_memory_manager.cc
@@ -347,14 +347,17 @@ SharedMemoryManager::GetMemoryInfo(
             .c_str());
   }
 
+  // populate 'shm_mapped_addr'
   size_t max_offset = 0;
   if (it->second->kind_ == TRITONSERVER_MEMORY_CPU) {
     *shm_mapped_addr = (void*)((uint8_t*)it->second->mapped_addr_ +
                                it->second->offset_ + offset);
-    max_offset = it->second->offset_ + it->second->byte_size_;
+    max_offset = it->second->offset_;
   } else {
     *shm_mapped_addr = (void*)((uint8_t*)it->second->mapped_addr_ + offset);
-    max_offset = it->second->byte_size_;
+  }
+  if (it->second->byte_size_ > 0) {
+    max_offset += it->second->byte_size_ - 1;
   }
   if (offset > max_offset) {
     *shm_mapped_addr = nullptr;