tritonserver-rs 0.4.1

// Copyright 2020-2022, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
//
// Redistribution and use in source and binary forms, with or without
// modification, are permitted provided that the following conditions
// are met:
//  * Redistributions of source code must retain the above copyright
//    notice, this list of conditions and the following disclaimer.
//  * Redistributions in binary form must reproduce the above copyright
//    notice, this list of conditions and the following disclaimer in the
//    documentation and/or other materials provided with the distribution.
//  * Neither the name of NVIDIA CORPORATION nor the names of its
//    contributors may be used to endorse or promote products derived
//    from this software without specific prior written permission.
//
// THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS ``AS IS'' AND ANY
// EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
// IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR
// PURPOSE ARE DISCLAIMED.  IN NO EVENT SHALL THE COPYRIGHT OWNER OR
// CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,
// EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,
// PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR
// PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY
// OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
// (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
// OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
#include "memory.h"

#include <cuda_runtime_api.h>

#include "cuda_block_manager.h"
#include "cuda_memory_manager.h"
#include "cuda_utils.h"
#include "gtest/gtest.h"
#include "pinned_memory_manager.h"
#include "triton/core/tritonserver.h"

namespace tc = triton::core;

namespace {

#define CHECK_POINTER_ATTRIBUTES(ptr__, type__, device__)                   \
  do {                                                                      \
    cudaPointerAttributes attr;                                             \
    auto cuerr = cudaPointerGetAttributes(&attr, ptr__);                    \
    ASSERT_TRUE(cuerr == cudaSuccess)                                       \
        << "Failed to get CUDA pointer attributes: "                        \
        << cudaGetErrorString(cuerr);                                       \
    EXPECT_TRUE(attr.type == type__)                                        \
        << "Expect pointer with type " << type__ << ", got: " << attr.type; \
    if (attr.type == cudaMemoryTypeDevice) {                                \
      EXPECT_TRUE(attr.device == device__)                                  \
          << "Expect allocation on CUDA device " << device__                \
          << ", got: " << attr.device;                                      \
    }                                                                       \
  } while (false)

// Wrapper of CudaMemoryManager class to expose Reset() for unit testing
class TestingCudaMemoryManager : public tc::CudaMemoryManager {
 public:
  static void Reset() { CudaMemoryManager::Reset(); }
};

class CudaMemoryManagerTest : public ::testing::Test {
 protected:
  void SetUp() override
  {
    // Default memory manager options
    options_.min_supported_compute_capability_ = 6.0;
    options_.memory_pool_byte_size_ = {{0, 1 << 10}};
  }

  void TearDown() override { TestingCudaMemoryManager::Reset(); }

  tc::CudaMemoryManager::Options options_;
};

TEST_F(CudaMemoryManagerTest, InitOOM)
{
  // Set to reserve too much memory
  double cc = 6.0;
  std::map<int, uint64_t> s{{0, uint64_t(1) << 40 /* 1024 GB */}};
  const tc::CudaMemoryManager::Options options{cc, s};
  auto status = tc::CudaMemoryManager::Create(options);
  EXPECT_FALSE(status.IsOk()) << "Expect creation error";
}

TEST_F(CudaMemoryManagerTest, InitSuccess)
{
  double cc = 6.0;
  std::map<int, uint64_t> s{{0, 1 << 10 /* 1024 bytes */}};
  const tc::CudaMemoryManager::Options options{cc, s};
  auto status = tc::CudaMemoryManager::Create(options);
  EXPECT_TRUE(status.IsOk()) << status.Message();
}

TEST_F(CudaMemoryManagerTest, InitNoDeviceConfig)
{
  double cc = 6.0;
  std::map<int, uint64_t> s;
  const tc::CudaMemoryManager::Options options{cc, s};
  auto status = tc::CudaMemoryManager::Create(options);
  EXPECT_TRUE(status.IsOk()) << status.Message();

  void* ptr = nullptr;
  status = tc::CudaMemoryManager::Alloc(&ptr, 1, 0);
  ASSERT_FALSE(status.IsOk()) << "Unexpected successful allocation";
}

TEST_F(CudaMemoryManagerTest, InitZeroByte)
{
  double cc = 6.0;
  std::map<int, uint64_t> s{{0, 0}};
  const tc::CudaMemoryManager::Options options{cc, s};
  auto status = tc::CudaMemoryManager::Create(options);
  EXPECT_TRUE(status.IsOk()) << status.Message();

  void* ptr = nullptr;
  status = tc::CudaMemoryManager::Alloc(&ptr, 1, 0);
  ASSERT_FALSE(status.IsOk()) << "Unexpected successful allocation";
}

TEST_F(CudaMemoryManagerTest, AllocSuccess)
{
  auto status = tc::CudaMemoryManager::Create(options_);
  ASSERT_TRUE(status.IsOk()) << status.Message();

  void* ptr = nullptr;
  status = tc::CudaMemoryManager::Alloc(&ptr, 1024, 0);
  ASSERT_TRUE(status.IsOk()) << status.Message();
  ASSERT_TRUE(ptr) << "Expect pointer to allocated buffer";
  // check if returned pointer is CUDA pointer
  CHECK_POINTER_ATTRIBUTES(ptr, cudaMemoryTypeDevice, 0);
}

TEST_F(CudaMemoryManagerTest, AllocFail)
{
  auto status = tc::CudaMemoryManager::Create(options_);
  ASSERT_TRUE(status.IsOk()) << status.Message();

  void* ptr = nullptr;
  status = tc::CudaMemoryManager::Alloc(&ptr, 2048, 0);
  ASSERT_FALSE(status.IsOk()) << "Unexpected successful allocation";
}

TEST_F(CudaMemoryManagerTest, MultipleAlloc)
{
  auto status = tc::CudaMemoryManager::Create(options_);
  ASSERT_TRUE(status.IsOk()) << status.Message();

  void* first_ptr = nullptr;
  status = tc::CudaMemoryManager::Alloc(&first_ptr, 600, 0);
  ASSERT_TRUE(status.IsOk()) << status.Message();
  ASSERT_TRUE(first_ptr) << "Expect pointer to allocated buffer";
  // check if returned pointer is CUDA pointer
  CHECK_POINTER_ATTRIBUTES(first_ptr, cudaMemoryTypeDevice, 0);

  // 512 + 600 > 1024
  void* second_ptr = nullptr;
  status = tc::CudaMemoryManager::Alloc(&second_ptr, 512, 0);
  ASSERT_FALSE(status.IsOk()) << "Unexpected successful allocation";

  // Free the first pointer and retry the second one
  status = tc::CudaMemoryManager::Free(first_ptr, 0);
  EXPECT_TRUE(status.IsOk()) << status.Message();
  status = tc::CudaMemoryManager::Alloc(&second_ptr, 512, 0);
  ASSERT_TRUE(status.IsOk()) << status.Message();
  ASSERT_TRUE(second_ptr) << "Expect pointer to allocated buffer";
  // check if returned pointer is CUDA pointer
  CHECK_POINTER_ATTRIBUTES(second_ptr, cudaMemoryTypeDevice, 0);
}

TEST_F(CudaMemoryManagerTest, MultipleDevice)
{
  std::set<int> supported_gpus;
  auto status = tc::GetSupportedGPUs(
      &supported_gpus, options_.min_supported_compute_capability_);
  ASSERT_TRUE(status.IsOk()) << status.Message();
  ASSERT_GE(supported_gpus.size(), size_t(2))
      << "Test requires at least two supported CUDA devices";

  {
    double cc = 6.0;
    std::map<int, uint64_t> s;
    // Only enough memory is only reserved in one of the devices
    s[*supported_gpus.begin()] = 32;
    s[*(++supported_gpus.begin())] = 1024;
    const tc::CudaMemoryManager::Options options{cc, s};
    status = tc::CudaMemoryManager::Create(options);
    ASSERT_TRUE(status.IsOk()) << status.Message();
  }

  void* ptr = nullptr;
  // Allocation on small device
  int small_device = *supported_gpus.begin();
  status = tc::CudaMemoryManager::Alloc(&ptr, 1024, small_device);
  ASSERT_FALSE(status.IsOk()) << "Unexpected successful allocation";

  // Allocation on large device
  int large_device = *(++supported_gpus.begin());
  status = tc::CudaMemoryManager::Alloc(&ptr, 1024, large_device);
  ASSERT_TRUE(status.IsOk()) << status.Message();
  ASSERT_TRUE(ptr) << "Expect pointer to allocated buffer";
  // check if returned pointer is CUDA pointer
  CHECK_POINTER_ATTRIBUTES(ptr, cudaMemoryTypeDevice, large_device);

  // Free allocation ...
  status = tc::CudaMemoryManager::Free(ptr, small_device);
  EXPECT_FALSE(status.IsOk()) << "Unexpected deallocation on wrong device";
  status = tc::CudaMemoryManager::Free(ptr, large_device);
  EXPECT_TRUE(status.IsOk()) << status.Message();
}

class AllocatedMemoryTest : public ::testing::Test {
 protected:
  // Per-test-suite set-up.
  static void SetUpTestSuite()
  {
    // Pinned memory manager
    {
      tc::PinnedMemoryManager::Options options{1024};
      auto status = tc::PinnedMemoryManager::Create(options);
      ASSERT_TRUE(status.IsOk()) << status.Message();
    }
  }

  // Set up CUDA memory manager per test for special fallback case
  void SetUp() override
  {
    tc::CudaMemoryManager::Options options{6.0, {{0, 1 << 10}}};
    auto status = tc::CudaMemoryManager::Create(options);
    ASSERT_TRUE(status.IsOk()) << status.Message();
  }

  void TearDown() override { TestingCudaMemoryManager::Reset(); }
};

TEST_F(AllocatedMemoryTest, AllocGPU)
{
  size_t expect_size = 512, actual_size;
  TRITONSERVER_MemoryType expect_type = TRITONSERVER_MEMORY_GPU, actual_type;
  int64_t expect_id = 0, actual_id;
  tc::AllocatedMemory memory(expect_size, expect_type, expect_id);

  auto ptr = memory.BufferAt(0, &actual_size, &actual_type, &actual_id);
  EXPECT_EQ(expect_size, actual_size)
      << "Expect size: " << expect_size << ", got: " << actual_size;
  EXPECT_EQ(expect_type, actual_type)
      << "Expect type: " << expect_type << ", got: " << actual_type;
  EXPECT_EQ(expect_id, actual_id)
      << "Expect id: " << expect_id << ", got: " << actual_id;

  // Sanity check on the pointer property
  CHECK_POINTER_ATTRIBUTES(ptr, cudaMemoryTypeDevice, expect_id);
}

TEST_F(AllocatedMemoryTest, AllocPinned)
{
  size_t expect_size = 512, actual_size;
  TRITONSERVER_MemoryType expect_type = TRITONSERVER_MEMORY_CPU_PINNED,
                          actual_type;
  int64_t expect_id = 0, actual_id;
  tc::AllocatedMemory memory(expect_size, expect_type, expect_id);

  auto ptr = memory.BufferAt(0, &actual_size, &actual_type, &actual_id);
  EXPECT_EQ(expect_size, actual_size)
      << "Expect size: " << expect_size << ", got: " << actual_size;
  EXPECT_EQ(expect_type, actual_type)
      << "Expect type: " << expect_type << ", got: " << actual_type;
  EXPECT_EQ(expect_id, actual_id)
      << "Expect id: " << expect_id << ", got: " << actual_id;

  // Sanity check on the pointer property
  CHECK_POINTER_ATTRIBUTES(ptr, cudaMemoryTypeHost, expect_id);
}

TEST_F(AllocatedMemoryTest, AllocFallback)
{
  // Each allocation uses half of the target reserved memory
  size_t expect_size = 600, actual_size;
  TRITONSERVER_MemoryType expect_type = TRITONSERVER_MEMORY_GPU, actual_type;
  int64_t expect_id = 0, actual_id;

  // First allocation
  tc::AllocatedMemory cuda_memory(expect_size, expect_type, expect_id);

  auto ptr = cuda_memory.BufferAt(0, &actual_size, &actual_type, &actual_id);
  EXPECT_EQ(expect_size, actual_size)
      << "Expect size: " << expect_size << ", got: " << actual_size;
  EXPECT_EQ(expect_type, actual_type)
      << "Expect type: " << expect_type << ", got: " << actual_type;
  EXPECT_EQ(expect_id, actual_id)
      << "Expect id: " << expect_id << ", got: " << actual_id;

  // Sanity check on the pointer property
  CHECK_POINTER_ATTRIBUTES(ptr, cudaMemoryTypeDevice, expect_id);

  // Second allocation, should trigger fallback from CUDA -> pinned memory
  tc::AllocatedMemory pinned_memory(expect_size, expect_type, expect_id);

  ptr = pinned_memory.BufferAt(0, &actual_size, &actual_type, &actual_id);
  EXPECT_EQ(expect_size, actual_size)
      << "Expect size: " << expect_size << ", got: " << actual_size;
  EXPECT_EQ(TRITONSERVER_MEMORY_CPU_PINNED, actual_type)
      << "Expect type: " << TRITONSERVER_MEMORY_CPU_PINNED
      << ", got: " << actual_type;

  // Sanity check on the pointer property
  CHECK_POINTER_ATTRIBUTES(ptr, cudaMemoryTypeHost, expect_id);

  // Third allocation, CUDA -> pinned -> non-pinned
  tc::AllocatedMemory system_memory(expect_size, expect_type, expect_id);

  ptr = system_memory.BufferAt(0, &actual_size, &actual_type, &actual_id);
  EXPECT_EQ(expect_size, actual_size)
      << "Expect size: " << expect_size << ", got: " << actual_size;
  EXPECT_EQ(TRITONSERVER_MEMORY_CPU, actual_type)
      << "Expect type: " << TRITONSERVER_MEMORY_CPU_PINNED
      << ", got: " << actual_type;

  // Sanity check on the pointer property
  CHECK_POINTER_ATTRIBUTES(ptr, cudaMemoryTypeUnregistered, expect_id);
}

TEST_F(AllocatedMemoryTest, AllocFallbackNoCuda)
{
  // Test fallback in the case where CUDA memory manager is not properly created
  TestingCudaMemoryManager::Reset();

  size_t expect_size = 600, actual_size;
  TRITONSERVER_MemoryType expect_type = TRITONSERVER_MEMORY_GPU, actual_type;
  int64_t expect_id = 0, actual_id;

  // CUDA memory allocation should trigger fallback to allocate pinned memory
  tc::AllocatedMemory pinned_memory(expect_size, expect_type, expect_id);

  auto ptr = pinned_memory.BufferAt(0, &actual_size, &actual_type, &actual_id);
  EXPECT_EQ(expect_size, actual_size)
      << "Expect size: " << expect_size << ", got: " << actual_size;
  EXPECT_EQ(TRITONSERVER_MEMORY_CPU_PINNED, actual_type)
      << "Expect type: " << TRITONSERVER_MEMORY_CPU_PINNED
      << ", got: " << actual_type;

  // Sanity check on the pointer property
  CHECK_POINTER_ATTRIBUTES(ptr, cudaMemoryTypeHost, expect_id);
}

TEST_F(AllocatedMemoryTest, Release)
{
  // Similar to above, but verify that the memory will be released once
  // out of scope
  // Each allocation uses half of the target reserved memory
  size_t expect_size = 600, actual_size;
  TRITONSERVER_MemoryType expect_type = TRITONSERVER_MEMORY_GPU, actual_type;
  int64_t expect_id = 0, actual_id;

  {
    // First allocation
    tc::AllocatedMemory cuda_memory(expect_size, expect_type, expect_id);

    auto ptr = cuda_memory.BufferAt(0, &actual_size, &actual_type, &actual_id);
    EXPECT_EQ(expect_size, actual_size)
        << "Expect size: " << expect_size << ", got: " << actual_size;
    EXPECT_EQ(expect_type, actual_type)
        << "Expect type: " << expect_type << ", got: " << actual_type;
    EXPECT_EQ(expect_id, actual_id)
        << "Expect id: " << expect_id << ", got: " << actual_id;

    // Sanity check on the pointer property
    CHECK_POINTER_ATTRIBUTES(ptr, cudaMemoryTypeDevice, expect_id);

    // Second allocation, should trigger fallback from CUDA -> pinned memory
    tc::AllocatedMemory pinned_memory(expect_size, expect_type, expect_id);

    ptr = pinned_memory.BufferAt(0, &actual_size, &actual_type, &actual_id);
    EXPECT_EQ(expect_size, actual_size)
        << "Expect size: " << expect_size << ", got: " << actual_size;
    EXPECT_EQ(TRITONSERVER_MEMORY_CPU_PINNED, actual_type)
        << "Expect type: " << TRITONSERVER_MEMORY_CPU_PINNED
        << ", got: " << actual_type;

    // Sanity check on the pointer property
    CHECK_POINTER_ATTRIBUTES(ptr, cudaMemoryTypeHost, expect_id);
  }

  // Third allocation, should not trigger fallback
  tc::AllocatedMemory memory(expect_size, expect_type, expect_id);

  auto ptr = memory.BufferAt(0, &actual_size, &actual_type, &actual_id);
  EXPECT_EQ(expect_size, actual_size)
      << "Expect size: " << expect_size << ", got: " << actual_size;
  EXPECT_EQ(expect_type, actual_type)
      << "Expect type: " << expect_type << ", got: " << actual_type;

  // Sanity check on the pointer property
  CHECK_POINTER_ATTRIBUTES(ptr, cudaMemoryTypeDevice, expect_id);
}

class GrowableMemoryTest : public ::testing::Test {
 protected:
  // Per-test-suite set-up.
  static void SetUpTestSuite()
  {
    auto status = tc::CudaBlockManager::Create(6.0);
    EXPECT_TRUE(status.IsOk()) << status.Message();
  }

  void TearDown() override { tc::CudaBlockManager::Reset(); }
};

TEST_F(GrowableMemoryTest, AllocGPU)
{
  std::unique_ptr<tc::GrowableMemory> growable_memory;
  size_t block_size = tc::CudaBlockManager::BlockSize();
  auto status = tc::GrowableMemory::Create(
      growable_memory, block_size, TRITONSERVER_MEMORY_GPU, 0, block_size * 20);
  EXPECT_TRUE(status.IsOk()) << status.Message();

  auto blocks = growable_memory->GetAllocation()->Blocks();

  TRITONSERVER_MemoryType memory_type;
  int64_t memory_type_id;
  char* buffer = growable_memory->MutableBuffer(&memory_type, &memory_type_id);
  EXPECT_EQ(memory_type, TRITONSERVER_MEMORY_GPU);
  EXPECT_EQ(memory_type_id, 0);
  CHECK_POINTER_ATTRIBUTES(buffer, cudaMemoryTypeDevice, 0);
  // Make sure the memory is destructed without issues.
  growable_memory.reset(nullptr);

  // Allocate new memory with the same size.
  status = tc::GrowableMemory::Create(
      growable_memory, block_size, TRITONSERVER_MEMORY_GPU, 0, block_size * 20);
  buffer = growable_memory->MutableBuffer(&memory_type, &memory_type_id);
  EXPECT_EQ(memory_type, TRITONSERVER_MEMORY_GPU);
  EXPECT_EQ(memory_type_id, 0);
  CHECK_POINTER_ATTRIBUTES(buffer, cudaMemoryTypeDevice, 0);
  EXPECT_TRUE(status.IsOk()) << status.Message();
  auto new_blocks = growable_memory->GetAllocation()->Blocks();
  ASSERT_EQ(blocks[0], new_blocks[0])
      << "Expected the same blocks found different blocks";

  // Grow the memory
  status = growable_memory->Resize(block_size * 3);
  buffer = growable_memory->MutableBuffer(&memory_type, &memory_type_id);
  EXPECT_EQ(memory_type, TRITONSERVER_MEMORY_GPU);
  EXPECT_EQ(memory_type_id, 0);
  CHECK_POINTER_ATTRIBUTES(buffer, cudaMemoryTypeDevice, 0);
  EXPECT_TRUE(status.IsOk()) << status.Message();
  ASSERT_EQ(growable_memory->GetAllocation()->Blocks().size(), 3)
      << "expected 3 blocks.";

  // Try to grow the memory beyond the virtual address size
  status = growable_memory->Resize(block_size * 21);
  EXPECT_FALSE(status.IsOk()) << status.Message();

  // Try to allocate memory larger than the block size
  status = tc::GrowableMemory::Create(
      growable_memory, block_size * 21, TRITONSERVER_MEMORY_GPU, 0,
      block_size * 20);
  EXPECT_FALSE(status.IsOk()) << status.Message();

  // Request a CPU memory and make sure you get GPU memory back
  status = tc::GrowableMemory::Create(
      growable_memory, block_size, TRITONSERVER_MEMORY_CPU, 0, block_size * 20);
  buffer = growable_memory->MutableBuffer(&memory_type, &memory_type_id);
  EXPECT_EQ(memory_type, TRITONSERVER_MEMORY_GPU);
  EXPECT_EQ(memory_type_id, 0);
}

}  // namespace

int
main(int argc, char** argv)
{
  ::testing::InitGoogleTest(&argc, argv);
  return RUN_ALL_TESTS();
}