Enhance bound check for shm offset

[kthui]

Enhance shared memory bound checking during inference to make sure the offset supplied by the user cannot go beyond the shared memory region when the region was registered.

[GuanLuo]

Shouldn't you also check the byte size of the read / write as well? Couldn't the same issue happen if one specifies offset = (max_offset - 1) and instructs Triton to write > 1 bytes?

[kthui]

Shouldn't you also check the byte size of the read / write as well? Couldn't the same issue happen if one specifies offset = (max_offset - 1) and instructs Triton to write > 1 bytes?

This scope is exactly the next ticket after this one. We can address it in a different PR, unless you think both fixes must be committed at the same time.

[nv-kmcgill53]

This is increasing the scope by a little bit but could you provide an improvement to the shared memory manager where we check that the offset provided by the user is a multiple of the host's page size? You could probably test for this if mmap returns an error.

const size_t page_size = sysconf(_SC_PAGE_SIZE);
if ( (offset % page_size ) != 0) {
      return TRITONSERVER_ErrorNew(
        TRITONSERVER_ERROR_INTERNAL, std::string(
                                         "Bad argument: offset, " + std::to_string(offest) + 
                                         ", is not a multiple of the host's page size (" + 
                                         std::to_string(page_size) + "): " +
                                         std::string(std::strerror(errno)))
                                         .c_str());
}

[kthui]

@nv-kmcgill53 if I understand correctly, you are talking about an improved bound check when registering a shm region, not when inferencing with a shm region? The next next ticket after this one is about improving checks when registering a shm region. Maybe we can add this check on a separate PR, unless you think the shm registration bound check improvement and this one must be committed together?

[nnshah1]

Small suggestion - what if we refactored to set the offsets, check them, and then only set the shm variable? I think one thing was a little odd was seeing the address set and then the offsets checked.

By setting the variables you can also keep a single place to set the address - just a thought.

size_t max_offset = 0;
 
if (it->second->byte_size_ > 0) {
    max_offset = it->second->byte_size_ - 1;
  }
  if (it->second->kind_ == TRITONSERVER_MEMORY_CPU) {
     offset += it->second->offset;
     max_offset += it -> second-> offset;
  } 
 
  if (offset > max_offset) {
    *shm_mapped_addr = nullptr;
    return TRITONSERVER_ErrorNew(
        TRITONSERVER_ERROR_INVALID_ARG,
        std::string("Invalid offset for shared memory region: '" + name + "'")
            .c_str());
  }
   *shm_mapped_addr = (void*)((uint8_t*)it->second->mapped_addr_ + offset);

[kthui]

Sure. Separate logic between computation and validation

[nv-kmcgill53]

@nv-kmcgill53 if I understand correctly, you are talking about an improved bound check when registering a shm region, not when inferencing with a shm region? The next next ticket after this one is about improving checks when registering a shm region. Maybe we can add this check on a separate PR, unless you think the shm registration bound check improvement and this one must be committed together?

Yes, if there is other work slated for improvement on the registering of shared memory then, yes, that change should go in there.

[kthui]

@nv-kmcgill53 if I understand correctly, you are talking about an improved bound check when registering a shm region, not when inferencing with a shm region? The next next ticket after this one is about improving checks when registering a shm region. Maybe we can add this check on a separate PR, unless you think the shm registration bound check improvement and this one must be committed together?

Yes, if there is other work slated for improvement on the registering of shared memory then, yes, that change should go in there.

Sure. Linked this conversation in that ticket.

[rmccorm4]

LGTM assuming the byte size check will be in a follow up