src/model_loader.cc - main - rnekrasov/python_backend

src/model_loader.cc
267 строк10 KB
Kris Hung
Model loading API for Python BLS (#262)
29 июн 2023, 23:51
29 июн 2023, 23:51c48fd37
// Copyright 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 "model_loader.h"

#ifdef TRITON_PB_STUB
#include "pb_stub.h"
#endif

namespace triton { namespace backend { namespace python {

void
ModelLoader::SaveToSharedMemory(std::unique_ptr<SharedMemoryManager>& shm_pool)
{
  AllocatedSharedMemory<ModelLoaderRequestShm> model_loader_req_shm =
      shm_pool->Construct<ModelLoaderRequestShm>();
  model_loader_req_shm_ptr_ = model_loader_req_shm.data_.get();

  std::unique_ptr<PbString> name_shm = PbString::Create(shm_pool, name_);
  std::unique_ptr<PbString> version_shm = PbString::Create(shm_pool, version_);
  std::unique_ptr<PbString> config_shm = PbString::Create(shm_pool, config_);
  std::unique_ptr<PbMap> files_shm = PbMap::Create(shm_pool, files_);

  model_loader_req_shm_ptr_->name_shm_handle = name_shm->ShmHandle();
  model_loader_req_shm_ptr_->version_shm_handle = version_shm->ShmHandle();
  model_loader_req_shm_ptr_->config_shm_handle = config_shm->ShmHandle();
  model_loader_req_shm_ptr_->files_shm_handle = files_shm->ShmHandle();
  model_loader_req_shm_ptr_->unload_dependents = unload_dependents_;

  // Save the references to shared memory.
  model_loader_req_shm_ = std::move(model_loader_req_shm);
  name_shm_ = std::move(name_shm);
  version_shm_ = std::move(version_shm);
  config_shm_ = std::move(config_shm);
  files_shm_ = std::move(files_shm);

  shm_handle_ = model_loader_req_shm_.handle_;
}

std::unique_ptr<ModelLoader>
ModelLoader::LoadFromSharedMemory(
    std::unique_ptr<SharedMemoryManager>& shm_pool,
    bi::managed_external_buffer::handle_t handle)
{
  AllocatedSharedMemory<ModelLoaderRequestShm> model_loader_req_shm =
      shm_pool->Load<ModelLoaderRequestShm>(handle);
  ModelLoaderRequestShm* model_loader_req_shm_ptr =
      model_loader_req_shm.data_.get();

  std::unique_ptr<PbString> name_shm = PbString::LoadFromSharedMemory(
      shm_pool, model_loader_req_shm_ptr->name_shm_handle);
  std::unique_ptr<PbString> version_shm = PbString::LoadFromSharedMemory(
      shm_pool, model_loader_req_shm_ptr->version_shm_handle);
  std::unique_ptr<PbString> config_shm = PbString::LoadFromSharedMemory(
      shm_pool, model_loader_req_shm_ptr->config_shm_handle);
  std::unique_ptr<PbMap> files_shm = PbMap::LoadFromSharedMemory(
      shm_pool, model_loader_req_shm_ptr->files_shm_handle);

  return std::unique_ptr<ModelLoader>(new ModelLoader(
      model_loader_req_shm, name_shm, version_shm, config_shm, files_shm));
}

ModelLoader::ModelLoader(
    AllocatedSharedMemory<ModelLoaderRequestShm>& model_loader_req_shm,
    std::unique_ptr<PbString>& name_shm, std::unique_ptr<PbString>& version_shm,
    std::unique_ptr<PbString>& config_shm, std::unique_ptr<PbMap>& files_shm)
    : model_loader_req_shm_(std::move(model_loader_req_shm)),
      name_shm_(std::move(name_shm)), version_shm_(std::move(version_shm)),
      config_shm_(std::move(config_shm)), files_shm_(std::move(files_shm))
{
  model_loader_req_shm_ptr_ = model_loader_req_shm_.data_.get();
  name_ = name_shm_->String();
  version_ = version_shm_->String();
  config_ = config_shm_->String();
  files_ = files_shm_->UnorderedMap();
  unload_dependents_ = model_loader_req_shm_ptr_->unload_dependents;
}
#ifdef TRITON_PB_STUB
void
ModelLoader::SendLoadModelRequest()
{
  std::unique_ptr<Stub>& stub = Stub::GetOrCreateInstance();
  SaveToSharedMemory(stub->ShmPool());
  AllocatedSharedMemory<ModelLoaderMessage> model_loader_msg_shm;

  try {
    stub->SendMessage<ModelLoaderMessage>(
        model_loader_msg_shm, PYTHONSTUB_LoadModelRequest, shm_handle_);
  }
  catch (const PythonBackendException& pb_exception) {
    throw PythonBackendException(
        "Failed to load model: " + std::string(pb_exception.what()));
  }
}

void
ModelLoader::SendUnloadModelRequest()
{
  std::unique_ptr<Stub>& stub = Stub::GetOrCreateInstance();
  SaveToSharedMemory(stub->ShmPool());
  AllocatedSharedMemory<ModelLoaderMessage> model_loader_msg_shm;
  try {
    stub->SendMessage<ModelLoaderMessage>(
        model_loader_msg_shm, PYTHONSTUB_UnloadModelRequest, shm_handle_);
  }
  catch (const PythonBackendException& pb_exception) {
    throw PythonBackendException(
        "Failed to unload model: " + std::string(pb_exception.what()));
  }
}

bool
ModelLoader::SendModelReadinessRequest()
{
  std::unique_ptr<Stub>& stub = Stub::GetOrCreateInstance();
  SaveToSharedMemory(stub->ShmPool());
  ModelLoaderMessage* model_loader_msg = nullptr;
  AllocatedSharedMemory<ModelLoaderMessage> model_loader_msg_shm;
  try {
    stub->SendMessage<ModelLoaderMessage>(
        model_loader_msg_shm, PYTHONSTUB_ModelReadinessRequest, shm_handle_);
  }
  catch (const PythonBackendException& pb_exception) {
    throw PythonBackendException(
        "Failed to check model readiness: " + std::string(pb_exception.what()));
  }

  model_loader_msg = model_loader_msg_shm.data_.get();
  return model_loader_msg->is_model_ready;
}

void
LoadModel(
    const std::string& name, const std::string& config, const py::object& files)
{
  std::unordered_map<std::string, std::string> files_map;

  if (!files.is_none()) {
    if (!py::isinstance<py::dict>(files)) {
      throw PythonBackendException(
          "failed to load model '" + name +
          "', files should be a dictionary of file paths and file contents");
    }

    py::dict files_dict = py::cast<py::dict>(files);
    for (const auto& item : files_dict) {
      std::string key = py::cast<std::string>(item.first);
      py::bytes value = py::cast<py::bytes>(item.second);
      std::string content(value);
      files_map[key] = content;
    }
  }

  ModelLoader model_loader(name, config, files_map);
  model_loader.SendLoadModelRequest();
}

void
UnloadModel(const std::string& name, const bool unload_dependents)
{
  ModelLoader model_loader(name, unload_dependents);
  model_loader.SendUnloadModelRequest();
}

bool
IsModelReady(const std::string& name, const std::string& version)
{
  ModelLoader model_loader(name, version);
  return model_loader.SendModelReadinessRequest();
}
#else
void
ModelLoader::LoadModel(TRITONSERVER_Server* server)
{
  std::string path = "";
  std::string file_content = "";
  std::vector<const TRITONSERVER_Parameter*> const_params;
  if (!config_.empty()) {
    const_params.emplace_back(TRITONSERVER_ParameterNew(
        "config", TRITONSERVER_PARAMETER_STRING, config_.c_str()));
  }
  if (!files_.empty()) {
    for (auto& file : files_) {
      path = file.first;
      file_content = file.second;
      const_params.emplace_back(TRITONSERVER_ParameterBytesNew(
          path.c_str(), file_content.data(), file_content.size()));
    }
  }

  THROW_IF_TRITON_ERROR(TRITONSERVER_ServerLoadModelWithParameters(
      server, name_.c_str(), const_params.data(), const_params.size()));

  for (const auto param : const_params) {
    TRITONSERVER_ParameterDelete(const_cast<TRITONSERVER_Parameter*>(param));
  }
}

void
ModelLoader::UnloadModel(TRITONSERVER_Server* server)
{
  if (unload_dependents_) {
    THROW_IF_TRITON_ERROR(
        TRITONSERVER_ServerUnloadModelAndDependents(server, name_.c_str()));
  } else {
    THROW_IF_TRITON_ERROR(
        TRITONSERVER_ServerUnloadModel(server, name_.c_str()));
  }
}

bool
ModelLoader::IsModelReady(TRITONSERVER_Server* server)
{
  bool is_ready = false;
  int64_t model_version = GetModelVersionFromString(version_);
  THROW_IF_TRITON_ERROR(TRITONSERVER_ServerModelIsReady(
      server, name_.c_str(), model_version, &is_ready));
  return is_ready;
}

int64_t
ModelLoader::GetModelVersionFromString(const std::string& version_string)
{
  int64_t version = -1;
  if (!version_string.empty()) {
    try {
      version = std::stol(version_string);
    }
    catch (std::exception& e) {
      throw PythonBackendException(
          "failed to get model version from specified version string '" +
          version_string + "' (details: " + e.what() +
          "), version should be an integral value > 0");
    }

    if (version < 0) {
      throw PythonBackendException(
          "failed to get model version from specified version string '" +
          version_string + "', version should be an integral value > 0");
    }
  }
  return version;
}
#endif
}}}  // namespace triton::backend::python