Error reduction: aligned --> aligas, cmake option edit, intermediate …

…stub launcher changes
triton-inference-server · Sep 12, 2023 · d73bd1d · d73bd1d
1 parent c10352e
commit d73bd1d
Show file tree

Hide file tree

Showing 6 changed files with 58 additions and 21 deletions.
diff --git a/CMakeLists.txt b/CMakeLists.txt
@@ -93,6 +93,9 @@ FetchContent_Declare(
   GIT_TAG "v0.8"
   GIT_SHALLOW ON
 )
+# Option must be set off so WIN32 build does not break
+set(CMAKE_POLICY_DEFAULT_CMP0077 NEW)
+set(BUILD_MOCK OFF)
 FetchContent_MakeAvailable(dlpack)
 
 #

diff --git a/src/pb_env.h b/src/pb_env.h
@@ -32,7 +32,7 @@
 
 #ifdef WIN32
 #include <windows.h>
-#undef  PATH_MAX
+#undef PATH_MAX
 #define PATH_MAX MAX_PATH
 #endif
 namespace triton { namespace backend { namespace python {

diff --git a/src/python_be.h b/src/python_be.h
@@ -26,6 +26,9 @@
 
 #pragma once
 
+#include <sys/stat.h>
+#include <sys/types.h>
+
 #include <array>
 #include <atomic>
 #include <boost/asio.hpp>
@@ -48,8 +51,6 @@
 #include <regex>
 #include <sstream>
 #include <string>
-#include <sys/stat.h>
-#include <sys/types.h>
 #include <thread>
 #include <unordered_map>
 #include <vector>

diff --git a/src/shm_manager.h b/src/shm_manager.h
@@ -62,9 +62,9 @@ struct AllocatedSharedMemory {
 // info is placed in the beginning and the actual object is placed after that
 // (i.e. 4 plus the aligned address is not 16-bytes aligned). The aligned memory
 // is required by semaphore otherwise it may lead to SIGBUS error on ARM.
-struct AllocatedShmOwnership {
+struct alignas(16) AllocatedShmOwnership {
   uint32_t ref_count_;
-} __attribute__((aligned(16)));
+};
 
 class SharedMemoryManager {
  public:

diff --git a/src/stub_launcher.cc b/src/stub_launcher.cc
@@ -31,21 +31,21 @@
 namespace triton { namespace backend { namespace python {
 
 StubLauncher::StubLauncher(const std::string stub_process_kind)
-    : parent_pid_(0), stub_pid_(0), is_initialized_(false),
-      stub_process_kind_(stub_process_kind), model_instance_name_(""),
-      device_id_(0), kind_("")
+    : is_initialized_(false), stub_process_kind_(stub_process_kind),
+      model_instance_name_(""), device_id_(0), kind_("")
 
 {
+  InitializeProcessHandlers();
 }
 
 StubLauncher::StubLauncher(
     const std::string stub_process_kind, const std::string model_instance_name,
     const int32_t device_id, const std::string kind)
-    : parent_pid_(0), stub_pid_(0), is_initialized_(false),
-      stub_process_kind_(stub_process_kind),
+    : is_initialized_(false), stub_process_kind_(stub_process_kind),
       model_instance_name_(model_instance_name), device_id_(device_id),
       kind_(kind)
 {
+  InitializeProcessHandlers();
 }
 
 TRITONSERVER_Error*
@@ -306,11 +306,10 @@ StubLauncher::Launch()
 
       // If the model is not initialized, wait for the stub process to exit.
       if (!is_initialized_) {
-        int status;
         stub_message_queue_.reset();
         parent_message_queue_.reset();
         memory_manager_.reset();
-        waitpid(stub_pid_, &status, 0);
+        WaitForStubProcess();
       }
     });
 
@@ -334,10 +333,7 @@ StubLauncher::Launch()
       }
       catch (const PythonBackendException& ex) {
         // Need to kill the stub process first
-        kill(stub_pid_, SIGKILL);
-        int status;
-        waitpid(stub_pid_, &status, 0);
-        stub_pid_ = 0;
+        KillStubProcess();
         throw BackendModelException(
             TRITONSERVER_ErrorNew(TRITONSERVER_ERROR_INTERNAL, ex.what()));
       }
@@ -514,11 +510,11 @@ StubLauncher::TerminateStub()
       force_kill = true;
     }
 
-    int status;
     if (force_kill) {
-      kill(stub_pid_, SIGKILL);
+      KillStubProcess();
+    } else {
+      WaitForStubProcess();
     }
-    waitpid(stub_pid_, &status, 0);
   }
 
   // First destroy the IPCControl. This makes sure that IPCControl is
@@ -539,10 +535,16 @@ StubLauncher::ClearQueues()
 void
 StubLauncher::KillStubProcess()
 {
+#ifdef _WIN32
+  uint32_t exit_code;
+  TerminateProcess(stub_pid_.hProcess, exit_code);
+  WaitForStubProcess();
+  CloseHandle(stub_pid_.hProcess);
+#else
   kill(stub_pid_, SIGKILL);
-  int status;
-  waitpid(stub_pid_, &status, 0);
+  WaitForStubProcess();
   stub_pid_ = 0;
+#endif
 }
 
 TRITONSERVER_Error*
@@ -598,4 +600,28 @@ StubLauncher::ReceiveMessageFromStub(
 
   return nullptr;  // success
 }
+
+void
+StubLauncher::InitializeProcessHandlers()
+{
+#ifdef _WIN32
+  ZeroMemory(&startup_info_, sizeof(startup_info_));
+  startup_info_.cb = sizeof(startup_info_);
+  ZeroMemory(&stub_pid_, sizeof(stub_pid_));
+#else
+  parent_pid_ = 0;
+  stub_pid_ = 0;
+#endif
+}
+
+void
+StubLauncher::WaitForStubProcess()
+{
+#ifdef _WIN32
+  WaitForSingleObject(stub_pid_.hProcess, INFINITE);
+#else
+  int status;
+  waitpid(stub_pid_, &status, 0);
+#endif
+}
 }}};  // namespace triton::backend::python
diff --git a/src/stub_launcher.h b/src/stub_launcher.h
@@ -149,8 +149,15 @@ class StubLauncher {
   TRITONSERVER_Error* ReceiveMessageFromStub(
       bi::managed_external_buffer::handle_t& message);
 
+  // Wait for stub process
+  void WaitForStubProcess();
+
+  // Initialize pid / process_information handlers
+  void InitializeProcessHandlers();
+
  private:
 #ifdef _WIN32
+  STARTUPINFO startup_info_;
   PROCESS_INFORMATION parent_pid_;
   PROCESS_INFORMATION stub_pid_;
 #else