src/pb_utils.cc - main - rnekrasov/python_backend

src/pb_utils.cc
326 строк9 KB
Iman Tabrizian
Randomize Python backend shared memory region naming (#351)
08 апр 2024, 20:22
08 апр 2024, 20:224d42111
// Copyright 2021-2023, 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 "pb_utils.h"

#ifdef _WIN32
#include <windows.h>

#include <algorithm>
#else
#include <dlfcn.h>
#endif


#ifdef TRITON_ENABLE_GPU
#include <cuda.h>
#include <cuda_runtime_api.h>
#endif

namespace triton { namespace backend { namespace python {

#ifdef TRITON_ENABLE_GPU

CUDAHandler::CUDAHandler()
{
  dl_open_handle_ = LoadSharedObject("libcuda.so");

  // If libcuda.so is successfully opened, it must be able to find
  // "cuPointerGetAttribute", "cuGetErrorString", and
  // "cuDevicePrimaryCtxGetState" symbols.
  if (dl_open_handle_ != nullptr) {
    void* cu_pointer_get_attribute_fn = LocateSymbol("cuPointerGetAttribute");
    if (cu_pointer_get_attribute_fn == nullptr) {
      throw PythonBackendException(
          std::string("Failed to locate 'cuPointerGetAttribute'. Error: ") +
          LocateSymbolError());
    }
    *((void**)&cu_pointer_get_attribute_fn_) = cu_pointer_get_attribute_fn;

    void* cu_get_error_string_fn = LocateSymbol("cuGetErrorString");
    if (cu_get_error_string_fn == nullptr) {
      throw PythonBackendException(
          std::string("Failed to locate 'cuGetErrorString'. Error: ") +
          LocateSymbolError());
    }
    *((void**)&cu_get_error_string_fn_) = cu_get_error_string_fn;

    void* cu_init_fn = LocateSymbol("cuInit");
    if (cu_init_fn == nullptr) {
      throw PythonBackendException(
          std::string("Failed to locate 'cuInit'. Error: ") +
          LocateSymbolError());
    }
    *((void**)&cu_init_fn_) = cu_init_fn;

    void* cu_device_primary_ctx_get_state_fn =
        LocateSymbol("cuDevicePrimaryCtxGetState");
    if (cu_device_primary_ctx_get_state_fn == nullptr) {
      throw PythonBackendException(
          std::string(
              "Failed to locate 'cuDevicePrimaryCtxGetState'. Error: ") +
          LocateSymbolError());
    }
    *((void**)&cu_device_primary_ctx_get_state_fn_) =
        cu_device_primary_ctx_get_state_fn;

    // Initialize the driver API.
    CUresult cuda_err = (*cu_init_fn_)(0 /* flags */);
    if (cuda_err != CUDA_SUCCESS) {
      const char* error_string;
      (*cu_get_error_string_fn_)(cuda_err, &error_string);
      error_str_ = std::string("failed to call cuInit: ") + error_string;
      CloseLibrary();
      dl_open_handle_ = nullptr;
    }
  }
}

void
CUDAHandler::PointerGetAttribute(
    CUdeviceptr* start_address, CUpointer_attribute attribute,
    CUdeviceptr dev_ptr)
{
  CUresult cuda_err =
      (*cu_pointer_get_attribute_fn_)(start_address, attribute, dev_ptr);
  if (cuda_err != CUDA_SUCCESS) {
    const char* error_string;
    (*cu_get_error_string_fn_)(cuda_err, &error_string);
    throw PythonBackendException(
        std::string(
            "failed to get cuda pointer device attribute: " +
            std::string(error_string))
            .c_str());
  }
}

bool
CUDAHandler::IsAvailable()
{
  return dl_open_handle_ != nullptr;
}

void
CUDAHandler::OpenCudaHandle(
    int64_t memory_type_id, cudaIpcMemHandle_t* cuda_mem_handle,
    void** data_ptr)
{
  std::lock_guard<std::mutex> guard{mu_};
  ScopedSetDevice scoped_set_device(memory_type_id);

  cudaError_t err = cudaIpcOpenMemHandle(
      data_ptr, *cuda_mem_handle, cudaIpcMemLazyEnablePeerAccess);
  if (err != cudaSuccess) {
    throw PythonBackendException(
        std::string("Failed to open the cudaIpcHandle. error: ") +
        cudaGetErrorString(err));
  }
}

void
CUDAHandler::CloseCudaHandle(int64_t memory_type_id, void* data_ptr)
{
  std::lock_guard<std::mutex> guard{mu_};
  int current_device;

  // Save the previous device
  cudaError_t err = cudaGetDevice(&current_device);
  if (err != cudaSuccess) {
    throw PythonBackendException(
        std::string("Failed to get the current CUDA device. error: ") +
        cudaGetErrorString(err));
  }

  // Restore the previous device before returning from the function.
  ScopedSetDevice scoped_set_device(memory_type_id);
  err = cudaIpcCloseMemHandle(data_ptr);
  if (err != cudaSuccess) {
    throw PythonBackendException(
        std::string("Failed to close the cudaIpcHandle. error: ") +
        cudaGetErrorString(err));
  }
}

bool
CUDAHandler::HasPrimaryContext(int device)
{
  unsigned int ctx_flags;
  int ctx_is_active = 0;
  CUresult cuda_err = (*cu_device_primary_ctx_get_state_fn_)(
      device, &ctx_flags, &ctx_is_active);
  if (cuda_err != CUDA_SUCCESS) {
    const char* error_string;
    (*cu_get_error_string_fn_)(cuda_err, &error_string);
    throw PythonBackendException(
        std::string(
            "failed to get primary context state: " + std::string(error_string))
            .c_str());
  }

  return ctx_is_active == 1;
}

void
CUDAHandler::MaybeSetDevice(int device)
{
  if (HasPrimaryContext(device)) {
    cudaError_t err = cudaSetDevice(device);
    if (err != cudaSuccess) {
      throw PythonBackendException(
          std::string("Failed to set the CUDA device to ") +
          std::to_string(device) + ". error: " + cudaGetErrorString(err));
    }
  }
}


CUDAHandler::~CUDAHandler() noexcept(false)
{
  if (dl_open_handle_ != nullptr) {
    CloseLibrary();
  }
}

void*
CUDAHandler::LoadSharedObject(const char* filename)
{
#ifdef _WIN32
  // NOTE: 'nvcuda.dll' is a placeholder library. Apparently, this should be the
  // equivalent library for Windows, but need to verify.
  return LoadLibraryA("nvcuda.dll");
#else
  return dlopen("libcuda.so", RTLD_LAZY);
#endif
}

void*
CUDAHandler::LocateSymbol(const char* symbol)
{
#ifdef _WIN32
  return GetProcAddress(static_cast<HMODULE>(dl_open_handle_), symbol);
#else
  return dlsym(dl_open_handle_, symbol);
#endif
}


std::string
CUDAHandler::LocateSymbolError()
{
#ifdef _WIN32
  return std::to_string(GetLastError());
#else
  return dlerror();
#endif
}

void
CUDAHandler::CloseLibrary()
{
  bool successful = true;
#ifdef _WIN32
  successful = (FreeLibrary(static_cast<HMODULE>(dl_open_handle_)) != 0);
#else
  successful = (dlclose(dl_open_handle_) == 0);
#endif
  if (!successful) {
    throw PythonBackendException("Failed to close the cuda library handle.");
  }
}


ScopedSetDevice::ScopedSetDevice(int device)
{
  device_ = device;
  THROW_IF_CUDA_ERROR(cudaGetDevice(&current_device_));

  if (current_device_ != device_) {
    THROW_IF_CUDA_ERROR(cudaSetDevice(device_));
  }
}

ScopedSetDevice::~ScopedSetDevice()
{
  if (current_device_ != device_) {
    CUDAHandler& cuda_handler = CUDAHandler::getInstance();
    cuda_handler.MaybeSetDevice(current_device_);
  }
}

bool
IsUsingCUDAPool(
    std::unique_ptr<CUDAMemoryPoolManager>& cuda_pool, int64_t memory_type_id,
    void* data)
{
  CUDAHandler& cuda_api = CUDAHandler::getInstance();
  CUdeviceptr cuda_pool_address = 0;
  cuda_api.PointerGetAttribute(
      &cuda_pool_address, CU_POINTER_ATTRIBUTE_RANGE_START_ADDR,
      reinterpret_cast<CUdeviceptr>(data));

  return (
      cuda_pool->CUDAPoolAddress(memory_type_id) ==
      reinterpret_cast<void*>(cuda_pool_address));
}

#endif  // TRITON_ENABLE_GPU

// FIXME: [DLIS-6078]: We should not need this function. However, some paths are
// being retrieved from core that are not platform-agnostic.
void
SanitizePath(std::string& path)
{
  std::replace(path.begin(), path.end(), '/', '\\');
}

#ifndef TRITON_PB_STUB
std::shared_ptr<TRITONSERVER_Error*>
WrapTritonErrorInSharedPtr(TRITONSERVER_Error* error)
{
  std::shared_ptr<TRITONSERVER_Error*> response_error(
      new TRITONSERVER_Error*, [](TRITONSERVER_Error** error) {
        if (error != nullptr && *error != nullptr) {
          TRITONSERVER_ErrorDelete(*error);
        }

        if (error != nullptr) {
          delete error;
        }
      });
  *response_error = error;
  return response_error;
}
#endif  // NOT TRITON_PB_STUB

std::string
GenerateUUID()
{
  static boost::uuids::random_generator generator;
  boost::uuids::uuid uuid = generator();
  return boost::uuids::to_string(uuid);
}

}}}  // namespace triton::backend::python