From a418923ff9c676d332769485b2634639cd7b5e97 Mon Sep 17 00:00:00 2001 From: Matthew Wittwer Date: Mon, 30 Jun 2025 10:56:53 -0700 Subject: [PATCH 1/2] Added input validation to model load --- src/pb_utils.cc | 135 +++++++++++++++++++ src/pb_utils.h | 21 +++ src/stub_launcher.cc | 307 +++++++++++++++++++++++++++---------------- 3 files changed, 349 insertions(+), 114 deletions(-) diff --git a/src/pb_utils.cc b/src/pb_utils.cc index 809531b8..957bd6ce 100644 --- a/src/pb_utils.cc +++ b/src/pb_utils.cc @@ -26,12 +26,20 @@ #include "pb_utils.h" +#include +#include + #ifdef _WIN32 #include #include #else #include +#include +#endif + +#ifndef _WIN32 +extern char **environ; #endif @@ -315,6 +323,42 @@ WrapTritonErrorInSharedPtr(TRITONSERVER_Error* error) } #endif // NOT TRITON_PB_STUB +bool IsValidIdentifier(const std::string& input) { + if (input.empty()) { + return false; + } + + // Check for invalid characters + if (input.find_first_of(INVALID_CHARS) != std::string::npos) { + return false; + } + + return true; +} + +bool IsValidPath(const std::string& path) { + if (path.empty()) { + return false; + } + + // Must be absolute path + if (path[0] != '/') { + return false; + } + + return true; +} + +bool IsExecutableFile(const std::string& filepath) { + struct stat file_stat; + if (stat(filepath.c_str(), &file_stat) != 0) { + return false; + } + + // Check if it's a regular file and executable by owner + return S_ISREG(file_stat.st_mode) && (file_stat.st_mode & S_IXUSR); +} + std::string GenerateUUID() { @@ -323,4 +367,95 @@ GenerateUUID() return boost::uuids::to_string(uuid); } +// Helper function to parse environment variables from activation script +std::map +ParseActivationScript(const std::string& activate_path) { + std::map env_vars; + + // Read the current environment as baseline +#ifndef _WIN32 + if (environ != nullptr) { + for (char **env = environ; *env != nullptr; env++) { + std::string env_str(*env); + size_t eq_pos = env_str.find('='); + if (eq_pos != std::string::npos) { + std::string key = env_str.substr(0, eq_pos); + std::string value = env_str.substr(eq_pos + 1); + env_vars[key] = value; + } + } + } +#endif + + // Parse activation script for environment changes + std::ifstream activate_file(activate_path); + if (!activate_file.is_open()) { + return env_vars; // Return current environment if can't read activation script + } + + std::string line; + while (std::getline(activate_file, line)) { + // Look for export statements or direct assignments + if (line.find("export ") == 0) { + // Handle: export VAR=value + line = line.substr(7); // Remove "export " + } + + size_t eq_pos = line.find('='); + if (eq_pos != std::string::npos && line[0] != '#') { + std::string key = line.substr(0, eq_pos); + std::string value = line.substr(eq_pos + 1); + + // Remove quotes if present + if (value.size() >= 2 && + ((value[0] == '"' && value.back() == '"') || + (value[0] == '\'' && value.back() == '\''))) { + value = value.substr(1, value.size() - 2); + } + + // Handle variable substitution for common cases + if (value.find("$PATH") != std::string::npos) { + size_t pos = value.find("$PATH"); + value.replace(pos, 5, env_vars["PATH"]); + } + if (value.find("$LD_LIBRARY_PATH") != std::string::npos) { + size_t pos = value.find("$LD_LIBRARY_PATH"); + value.replace(pos, 16, env_vars["LD_LIBRARY_PATH"]); + } + + env_vars[key] = value; + } + } + + return env_vars; +} + +// Helper function to prepare environment array for execve +std::pair, std::vector> +PrepareEnvironment(const std::map& env_vars, + const std::string& additional_lib_path) { + + std::vector env_strings; + std::vector env_array; + + for (const auto& [key, value] : env_vars) { + std::string env_string; + if (key == "LD_LIBRARY_PATH" && !additional_lib_path.empty()) { + // Prepend the additional library path + env_string = key + "=" + additional_lib_path + ":" + value; + } else { + env_string = key + "=" + value; + } + env_strings.push_back(env_string); + } + + // Convert to char* array + for (auto& env_str : env_strings) { + env_array.push_back(const_cast(env_str.c_str())); + } + env_array.push_back(nullptr); + + return std::make_pair(std::move(env_strings), std::move(env_array)); +} + }}} // namespace triton::backend::python diff --git a/src/pb_utils.h b/src/pb_utils.h index 1306f375..65b24b1d 100644 --- a/src/pb_utils.h +++ b/src/pb_utils.h @@ -36,10 +36,12 @@ #include #include #include +#include #include #include #include #include +#include #include #include "pb_exception.h" @@ -341,6 +343,18 @@ bool IsUsingCUDAPool( // being retrieved from core that are not platform-agnostic. void SanitizePath(std::string& path); +// Invalid characters that are not allowed in user input +constexpr const char* INVALID_CHARS = ";|&$`<>()[]{}\\\"'*?~#!"; + +// Validate that an identifier (model name, region name, etc.) +bool IsValidIdentifier(const std::string& input); + +// Validate that a path exists and is absolute +bool IsValidPath(const std::string& path); + +// Check if a file exists and is executable +bool IsExecutableFile(const std::string& filepath); + #ifndef TRITON_PB_STUB std::shared_ptr WrapTritonErrorInSharedPtr( TRITONSERVER_Error* error); @@ -348,4 +362,11 @@ std::shared_ptr WrapTritonErrorInSharedPtr( std::string GenerateUUID(); +// Environment handling utilities for Python activation scripts +std::map ParseActivationScript(const std::string& activate_path); + +std::pair, std::vector> PrepareEnvironment( + const std::map& env_vars, + const std::string& additional_lib_path = ""); + }}} // namespace triton::backend::python diff --git a/src/stub_launcher.cc b/src/stub_launcher.cc index 828228e6..be3b47bb 100644 --- a/src/stub_launcher.cc +++ b/src/stub_launcher.cc @@ -29,11 +29,14 @@ #include #include "python_be.h" +#include "pb_utils.h" #ifdef _WIN32 #include // getpid() #endif +extern char **environ; + namespace triton { namespace backend { namespace python { StubLauncher::StubLauncher(const std::string stub_process_kind) @@ -337,10 +340,23 @@ StubLauncher::Launch() stub_name = model_instance_name_; } - const char* stub_args[4]; - stub_args[0] = "bash"; - stub_args[1] = "-c"; - stub_args[3] = nullptr; // Last argument must be nullptr + if (!IsValidIdentifier(stub_name)) { + return TRITONSERVER_ErrorNew( + TRITONSERVER_ERROR_INVALID_ARG, + "Invalid stub name: contains invalid characters"); + } + + if (!IsValidPath(model_path_)) { + return TRITONSERVER_ErrorNew( + TRITONSERVER_ERROR_INVALID_ARG, + "Invalid model path: contains invalid characters or not absolute"); + } + + if (!IsValidIdentifier(shm_region_name_)) { + return TRITONSERVER_ErrorNew( + TRITONSERVER_ERROR_INVALID_ARG, + "Invalid shared memory region name: contains invalid characters"); + } // Default Python backend stub std::string python_backend_stub = python_lib_ + "/triton_python_backend_stub"; @@ -353,50 +369,15 @@ StubLauncher::Launch() python_backend_stub = model_python_backend_stub; } - std::string bash_argument; - - // This shared memory variable indicates whether the stub process should - // revert the LD_LIBRARY_PATH changes to avoid shared library issues in - // executables and libraries. - ipc_control_->uses_env = false; - if (python_execution_env_ != "") { - std::stringstream ss; - - // Need to properly set the LD_LIBRARY_PATH so that Python environments - // using different python versions load properly. - ss << "source " << path_to_activate_ - << " && exec env LD_LIBRARY_PATH=" << path_to_libpython_ - << ":$LD_LIBRARY_PATH " << python_backend_stub << " " << model_path_ - << " " << shm_region_name_ << " " << shm_default_byte_size_ << " " - << shm_growth_byte_size_ << " " << parent_pid_ << " " << python_lib_ - << " " << ipc_control_handle_ << " " << stub_name << " " - << runtime_modeldir_; - ipc_control_->uses_env = true; - bash_argument = ss.str(); - } else { - std::stringstream ss; - ss << " exec " << python_backend_stub << " " << model_path_ << " " - << shm_region_name_ << " " << shm_default_byte_size_ << " " - << shm_growth_byte_size_ << " " << parent_pid_ << " " << python_lib_ - << " " << ipc_control_handle_ << " " << stub_name << " " - << runtime_modeldir_; - bash_argument = ss.str(); + // Validate the stub executable path + if (!IsValidPath(python_backend_stub)) { + return TRITONSERVER_ErrorNew( + TRITONSERVER_ERROR_INVALID_ARG, + "Invalid python backend stub path"); } - LOG_MESSAGE( - TRITONSERVER_LOG_VERBOSE, - (std::string("Starting Python backend stub: ") + bash_argument).c_str()); - - stub_args[2] = bash_argument.c_str(); - int stub_status_code = - system((python_backend_stub + "> /dev/null 2>&1").c_str()); - - // If running stub process without any arguments returns any status code, - // other than 1, it can indicate a permission issue as a result of - // downloading the stub process from a cloud object storage service. - if (WEXITSTATUS(stub_status_code) != 1) { - // Give the execute permission for the triton_python_backend_stub to the - // owner. + if (!IsExecutableFile(python_backend_stub)) { + // Give the execute permission for the triton_python_backend_stub to the owner. int error = chmod(python_backend_stub.c_str(), S_IXUSR); if (error != 0) { return TRITONSERVER_ErrorNew( @@ -409,79 +390,177 @@ StubLauncher::Launch() } } - pid_t pid = fork(); - if (pid < 0) { - return TRITONSERVER_ErrorNew( - TRITONSERVER_ERROR_INTERNAL, - "Failed to fork the stub process for auto-complete."); - } - if (pid == 0) { - // Replace this child process with the new stub process. - execvp("bash", (char**)stub_args); - // execvp() never return if succeeded. Otherwise, an error has occurred. - std::stringstream ss; - ss << "Failed to run python backend stub. Errno = " << errno << '\n' - << "Python backend stub path: " << python_backend_stub << '\n' - << "Shared Memory Region Name: " << shm_region_name_ << '\n' - << "Shared Memory Default Byte Size: " << shm_default_byte_size_ << '\n' - << "Shared Memory Growth Byte Size: " << shm_growth_byte_size_ << '\n'; - // Print the error message directly because the underlying mutexes in - // LOG_MESSAGE() could be forked when it is locked by other thread(s). - std::cerr << '\n' << ss.str() << '\n'; - // Terminate the child execution immediately to avoid any issues. - _Exit(1); - } else { - ScopedDefer _([&] { - // Push a dummy message to the message queue so that the stub - // process is notified that it can release the object stored in - // shared memory. - stub_message_queue_->Push(DUMMY_MESSAGE); - - // If the model is not initialized, wait for the stub process to exit. - if (!is_initialized_) { - stub_message_queue_.reset(); - parent_message_queue_.reset(); - memory_manager_.reset(); - WaitForStubProcess(); - } - }); - - stub_pid_ = pid; - - // The stub process would send two messages to the parent process during the - // initialization. - // 1. When the stub process's health monitoring thread has started. - // 2. When the initialization is fully completed and the Python model is - // loaded. - // - // The reason it is broken into two steps is that creation of the health - // monitoring thread may take longer which can make the server process think - // that the stub process is unhealthy and return early. Waiting until the - // health thread is spawn would prevent this issue. - parent_message_queue_->Pop(); - - if (stub_process_kind_ == "AUTOCOMPLETE_STUB") { - try { - AutocompleteStubProcess(); - } - catch (const PythonBackendException& ex) { - // Need to kill the stub process first - KillStubProcess(); - throw BackendModelException( - TRITONSERVER_ErrorNew(TRITONSERVER_ERROR_INTERNAL, ex.what())); - } - } else if (stub_process_kind_ == "MODEL_INSTANCE_STUB") { - RETURN_IF_ERROR(ModelInstanceStubProcess()); + // Prepare arguments for execution + std::vector arg_strings; + std::vector exec_args; + + // This shared memory variable indicates whether the stub process should + // revert the LD_LIBRARY_PATH changes to avoid shared library issues in + // executables and libraries. + ipc_control_->uses_env = false; + + if (python_execution_env_ != "") { + + // Validate Python environment paths + if (!IsValidPath(path_to_activate_) || !IsValidPath(path_to_libpython_)) { + return TRITONSERVER_ErrorNew( + TRITONSERVER_ERROR_INVALID_ARG, + "Invalid Python environment paths"); + } + + ipc_control_->uses_env = true; + + // Parse environment variables from activation script + std::map env_vars = ParseActivationScript(path_to_activate_); + + // Prepare environment with additional library path + auto [env_strings, custom_env] = PrepareEnvironment(env_vars, path_to_libpython_); + + // Set up arguments for direct execution + arg_strings.push_back(python_backend_stub); + arg_strings.push_back(model_path_); + arg_strings.push_back(shm_region_name_); + arg_strings.push_back(std::to_string(shm_default_byte_size_)); + arg_strings.push_back(std::to_string(shm_growth_byte_size_)); + arg_strings.push_back(std::to_string(parent_pid_)); + arg_strings.push_back(python_lib_); + arg_strings.push_back(std::to_string(ipc_control_handle_)); + arg_strings.push_back(stub_name); + arg_strings.push_back(runtime_modeldir_); + + // Convert strings to char* array for exec + for (const auto& arg : arg_strings) { + exec_args.push_back(arg.c_str()); + } + exec_args.push_back(nullptr); // exec requires null termination + + // Log the command being executed + std::ostringstream log_cmd; + for (size_t i = 0; i < arg_strings.size(); ++i) { + if (i > 0) log_cmd << " "; + log_cmd << "'" << arg_strings[i] << "'"; + } + LOG_MESSAGE( + TRITONSERVER_LOG_VERBOSE, + (std::string("Starting Python backend stub with custom environment: ") + log_cmd.str()).c_str()); + + pid_t pid = fork(); + if (pid < 0) { + return TRITONSERVER_ErrorNew( + TRITONSERVER_ERROR_INTERNAL, + "Failed to fork the stub process for auto-complete."); + } + if (pid == 0) { + // Replace this child process with the new stub process using custom environment + execve(python_backend_stub.c_str(), const_cast(exec_args.data()), custom_env.data()); + // execve() never returns if succeeded. Otherwise, an error has occurred. + std::stringstream ss; + ss << "Failed to run python backend stub with custom environment. Errno = " << errno << '\n' + << "Python backend stub path: " << python_backend_stub << '\n' + << "Activation script: " << path_to_activate_ << '\n' + << "Library path: " << path_to_libpython_ << '\n'; + std::cerr << '\n' << ss.str() << '\n'; + _Exit(1); } else { + stub_pid_ = pid; + } + + } else { + arg_strings.push_back(python_backend_stub); + arg_strings.push_back(model_path_); + arg_strings.push_back(shm_region_name_); + arg_strings.push_back(std::to_string(shm_default_byte_size_)); + arg_strings.push_back(std::to_string(shm_growth_byte_size_)); + arg_strings.push_back(std::to_string(parent_pid_)); + arg_strings.push_back(python_lib_); + arg_strings.push_back(std::to_string(ipc_control_handle_)); + arg_strings.push_back(stub_name); + arg_strings.push_back(runtime_modeldir_); + + // Convert strings to char* array for exec + for (const auto& arg : arg_strings) { + exec_args.push_back(arg.c_str()); + } + exec_args.push_back(nullptr); // exec requires null termination + + // Log the command being executed + std::ostringstream log_cmd; + for (size_t i = 0; i < arg_strings.size(); ++i) { + if (i > 0) log_cmd << " "; + log_cmd << "'" << arg_strings[i] << "'"; + } + LOG_MESSAGE( + TRITONSERVER_LOG_VERBOSE, + (std::string("Starting Python backend stub: ") + log_cmd.str()).c_str()); + + pid_t pid = fork(); + if (pid < 0) { return TRITONSERVER_ErrorNew( TRITONSERVER_ERROR_INTERNAL, - (std::string("Unknown stub_process_kind: ") + stub_process_kind_) - .c_str()); + "Failed to fork the stub process for auto-complete."); + } + if (pid == 0) { + // Replace this child process with the new stub process. + execv(python_backend_stub.c_str(), const_cast(exec_args.data())); + // execv() never returns if succeeded. Otherwise, an error has occurred. + std::stringstream ss; + ss << "Failed to run python backend stub. Errno = " << errno << '\n' + << "Python backend stub path: " << python_backend_stub << '\n'; + std::cerr << '\n' << ss.str() << '\n'; + _Exit(1); + } else { + stub_pid_ = pid; + } + } + + ScopedDefer _([&] { + // Push a dummy message to the message queue so that the stub + // process is notified that it can release the object stored in + // shared memory. + stub_message_queue_->Push(DUMMY_MESSAGE); + + // If the model is not initialized, wait for the stub process to exit. + if (!is_initialized_) { + stub_message_queue_.reset(); + parent_message_queue_.reset(); + memory_manager_.reset(); + WaitForStubProcess(); } + }); + + // The stub process would send two messages to the parent process during the + // initialization. + // 1. When the stub process's health monitoring thread has started. + // 2. When the initialization is fully completed and the Python model is + // loaded. + // + // The reason it is broken into two steps is that creation of the health + // monitoring thread may take longer which can make the server process think + // that the stub process is unhealthy and return early. Waiting until the + // health thread is spawn would prevent this issue. + parent_message_queue_->Pop(); - is_initialized_ = true; + if (stub_process_kind_ == "AUTOCOMPLETE_STUB") { + try { + AutocompleteStubProcess(); + } + catch (const PythonBackendException& ex) { + // Need to kill the stub process first + KillStubProcess(); + throw BackendModelException( + TRITONSERVER_ErrorNew(TRITONSERVER_ERROR_INTERNAL, ex.what())); + } + } else if (stub_process_kind_ == "MODEL_INSTANCE_STUB") { + RETURN_IF_ERROR(ModelInstanceStubProcess()); + } else { + return TRITONSERVER_ErrorNew( + TRITONSERVER_ERROR_INTERNAL, + (std::string("Unknown stub_process_kind: ") + stub_process_kind_) + .c_str()); } + is_initialized_ = true; + return nullptr; } @@ -895,6 +974,6 @@ StubLauncher::ShareCUDAMemoryPool( if (pb_exception.what() != std::string{""}) { throw pb_exception; } -} -#endif // TRITON_ENABLE_GPU + } + #endif // TRITON_ENABLE_GPU }}}; // namespace triton::backend::python From 988c0f30780051341dbbd115ef6caa1a07f170e9 Mon Sep 17 00:00:00 2001 From: Matthew Wittwer Date: Mon, 30 Jun 2025 20:42:12 -0700 Subject: [PATCH 2/2] clean up pre-commit errors --- src/pb_utils.cc | 73 +++++++++++++++++++++++++------------------- src/pb_utils.h | 5 +-- src/stub_launcher.cc | 70 ++++++++++++++++++++++++------------------ 3 files changed, 84 insertions(+), 64 deletions(-) diff --git a/src/pb_utils.cc b/src/pb_utils.cc index 957bd6ce..c96a0232 100644 --- a/src/pb_utils.cc +++ b/src/pb_utils.cc @@ -27,6 +27,7 @@ #include "pb_utils.h" #include + #include #ifdef _WIN32 @@ -39,7 +40,7 @@ #endif #ifndef _WIN32 -extern char **environ; +extern char** environ; #endif @@ -323,38 +324,44 @@ WrapTritonErrorInSharedPtr(TRITONSERVER_Error* error) } #endif // NOT TRITON_PB_STUB -bool IsValidIdentifier(const std::string& input) { +bool +IsValidIdentifier(const std::string& input) +{ if (input.empty()) { return false; } - + // Check for invalid characters if (input.find_first_of(INVALID_CHARS) != std::string::npos) { return false; } - + return true; } -bool IsValidPath(const std::string& path) { +bool +IsValidPath(const std::string& path) +{ if (path.empty()) { return false; } - + // Must be absolute path if (path[0] != '/') { return false; } - + return true; } -bool IsExecutableFile(const std::string& filepath) { +bool +IsExecutableFile(const std::string& filepath) +{ struct stat file_stat; if (stat(filepath.c_str(), &file_stat) != 0) { return false; } - + // Check if it's a regular file and executable by owner return S_ISREG(file_stat.st_mode) && (file_stat.st_mode & S_IXUSR); } @@ -368,14 +375,15 @@ GenerateUUID() } // Helper function to parse environment variables from activation script -std::map -ParseActivationScript(const std::string& activate_path) { +std::map +ParseActivationScript(const std::string& activate_path) +{ std::map env_vars; - + // Read the current environment as baseline #ifndef _WIN32 if (environ != nullptr) { - for (char **env = environ; *env != nullptr; env++) { + for (char** env = environ; *env != nullptr; env++) { std::string env_str(*env); size_t eq_pos = env_str.find('='); if (eq_pos != std::string::npos) { @@ -386,33 +394,33 @@ ParseActivationScript(const std::string& activate_path) { } } #endif - + // Parse activation script for environment changes std::ifstream activate_file(activate_path); if (!activate_file.is_open()) { - return env_vars; // Return current environment if can't read activation script + return env_vars; // Return current environment if can't read activation + // script } - + std::string line; while (std::getline(activate_file, line)) { // Look for export statements or direct assignments if (line.find("export ") == 0) { // Handle: export VAR=value - line = line.substr(7); // Remove "export " + line = line.substr(7); // Remove "export " } - + size_t eq_pos = line.find('='); if (eq_pos != std::string::npos && line[0] != '#') { std::string key = line.substr(0, eq_pos); std::string value = line.substr(eq_pos + 1); - + // Remove quotes if present - if (value.size() >= 2 && - ((value[0] == '"' && value.back() == '"') || - (value[0] == '\'' && value.back() == '\''))) { + if (value.size() >= 2 && ((value[0] == '"' && value.back() == '"') || + (value[0] == '\'' && value.back() == '\''))) { value = value.substr(1, value.size() - 2); } - + // Handle variable substitution for common cases if (value.find("$PATH") != std::string::npos) { size_t pos = value.find("$PATH"); @@ -422,22 +430,23 @@ ParseActivationScript(const std::string& activate_path) { size_t pos = value.find("$LD_LIBRARY_PATH"); value.replace(pos, 16, env_vars["LD_LIBRARY_PATH"]); } - + env_vars[key] = value; } } - + return env_vars; } // Helper function to prepare environment array for execve -std::pair, std::vector> -PrepareEnvironment(const std::map& env_vars, - const std::string& additional_lib_path) { - +std::pair, std::vector> +PrepareEnvironment( + const std::map& env_vars, + const std::string& additional_lib_path) +{ std::vector env_strings; std::vector env_array; - + for (const auto& [key, value] : env_vars) { std::string env_string; if (key == "LD_LIBRARY_PATH" && !additional_lib_path.empty()) { @@ -448,13 +457,13 @@ PrepareEnvironment(const std::map& env_vars, } env_strings.push_back(env_string); } - + // Convert to char* array for (auto& env_str : env_strings) { env_array.push_back(const_cast(env_str.c_str())); } env_array.push_back(nullptr); - + return std::make_pair(std::move(env_strings), std::move(env_array)); } diff --git a/src/pb_utils.h b/src/pb_utils.h index 65b24b1d..f8cbe1d0 100644 --- a/src/pb_utils.h +++ b/src/pb_utils.h @@ -363,10 +363,11 @@ std::shared_ptr WrapTritonErrorInSharedPtr( std::string GenerateUUID(); // Environment handling utilities for Python activation scripts -std::map ParseActivationScript(const std::string& activate_path); +std::map ParseActivationScript( + const std::string& activate_path); std::pair, std::vector> PrepareEnvironment( - const std::map& env_vars, + const std::map& env_vars, const std::string& additional_lib_path = ""); }}} // namespace triton::backend::python diff --git a/src/stub_launcher.cc b/src/stub_launcher.cc index be3b47bb..7cbecf77 100644 --- a/src/stub_launcher.cc +++ b/src/stub_launcher.cc @@ -28,14 +28,14 @@ #include -#include "python_be.h" #include "pb_utils.h" +#include "python_be.h" #ifdef _WIN32 #include // getpid() #endif -extern char **environ; +extern char** environ; namespace triton { namespace backend { namespace python { @@ -345,13 +345,13 @@ StubLauncher::Launch() TRITONSERVER_ERROR_INVALID_ARG, "Invalid stub name: contains invalid characters"); } - + if (!IsValidPath(model_path_)) { return TRITONSERVER_ErrorNew( TRITONSERVER_ERROR_INVALID_ARG, "Invalid model path: contains invalid characters or not absolute"); } - + if (!IsValidIdentifier(shm_region_name_)) { return TRITONSERVER_ErrorNew( TRITONSERVER_ERROR_INVALID_ARG, @@ -372,12 +372,12 @@ StubLauncher::Launch() // Validate the stub executable path if (!IsValidPath(python_backend_stub)) { return TRITONSERVER_ErrorNew( - TRITONSERVER_ERROR_INVALID_ARG, - "Invalid python backend stub path"); + TRITONSERVER_ERROR_INVALID_ARG, "Invalid python backend stub path"); } if (!IsExecutableFile(python_backend_stub)) { - // Give the execute permission for the triton_python_backend_stub to the owner. + // Give the execute permission for the triton_python_backend_stub to the + // owner. int error = chmod(python_backend_stub.c_str(), S_IXUSR); if (error != 0) { return TRITONSERVER_ErrorNew( @@ -398,24 +398,24 @@ StubLauncher::Launch() // revert the LD_LIBRARY_PATH changes to avoid shared library issues in // executables and libraries. ipc_control_->uses_env = false; - + if (python_execution_env_ != "") { - // Validate Python environment paths if (!IsValidPath(path_to_activate_) || !IsValidPath(path_to_libpython_)) { return TRITONSERVER_ErrorNew( - TRITONSERVER_ERROR_INVALID_ARG, - "Invalid Python environment paths"); + TRITONSERVER_ERROR_INVALID_ARG, "Invalid Python environment paths"); } - + ipc_control_->uses_env = true; - + // Parse environment variables from activation script - std::map env_vars = ParseActivationScript(path_to_activate_); - - // Prepare environment with additional library path - auto [env_strings, custom_env] = PrepareEnvironment(env_vars, path_to_libpython_); - + std::map env_vars = + ParseActivationScript(path_to_activate_); + + // Prepare environment with additional library path + auto [env_strings, custom_env] = + PrepareEnvironment(env_vars, path_to_libpython_); + // Set up arguments for direct execution arg_strings.push_back(python_backend_stub); arg_strings.push_back(model_path_); @@ -427,7 +427,7 @@ StubLauncher::Launch() arg_strings.push_back(std::to_string(ipc_control_handle_)); arg_strings.push_back(stub_name); arg_strings.push_back(runtime_modeldir_); - + // Convert strings to char* array for exec for (const auto& arg : arg_strings) { exec_args.push_back(arg.c_str()); @@ -437,12 +437,15 @@ StubLauncher::Launch() // Log the command being executed std::ostringstream log_cmd; for (size_t i = 0; i < arg_strings.size(); ++i) { - if (i > 0) log_cmd << " "; + if (i > 0) + log_cmd << " "; log_cmd << "'" << arg_strings[i] << "'"; } LOG_MESSAGE( TRITONSERVER_LOG_VERBOSE, - (std::string("Starting Python backend stub with custom environment: ") + log_cmd.str()).c_str()); + (std::string("Starting Python backend stub with custom environment: ") + + log_cmd.str()) + .c_str()); pid_t pid = fork(); if (pid < 0) { @@ -451,11 +454,16 @@ StubLauncher::Launch() "Failed to fork the stub process for auto-complete."); } if (pid == 0) { - // Replace this child process with the new stub process using custom environment - execve(python_backend_stub.c_str(), const_cast(exec_args.data()), custom_env.data()); + // Replace this child process with the new stub process using custom + // environment + execve( + python_backend_stub.c_str(), const_cast(exec_args.data()), + custom_env.data()); // execve() never returns if succeeded. Otherwise, an error has occurred. std::stringstream ss; - ss << "Failed to run python backend stub with custom environment. Errno = " << errno << '\n' + ss << "Failed to run python backend stub with custom environment. Errno " + "= " + << errno << '\n' << "Python backend stub path: " << python_backend_stub << '\n' << "Activation script: " << path_to_activate_ << '\n' << "Library path: " << path_to_libpython_ << '\n'; @@ -464,7 +472,7 @@ StubLauncher::Launch() } else { stub_pid_ = pid; } - + } else { arg_strings.push_back(python_backend_stub); arg_strings.push_back(model_path_); @@ -476,7 +484,7 @@ StubLauncher::Launch() arg_strings.push_back(std::to_string(ipc_control_handle_)); arg_strings.push_back(stub_name); arg_strings.push_back(runtime_modeldir_); - + // Convert strings to char* array for exec for (const auto& arg : arg_strings) { exec_args.push_back(arg.c_str()); @@ -486,12 +494,14 @@ StubLauncher::Launch() // Log the command being executed std::ostringstream log_cmd; for (size_t i = 0; i < arg_strings.size(); ++i) { - if (i > 0) log_cmd << " "; + if (i > 0) + log_cmd << " "; log_cmd << "'" << arg_strings[i] << "'"; } LOG_MESSAGE( TRITONSERVER_LOG_VERBOSE, - (std::string("Starting Python backend stub: ") + log_cmd.str()).c_str()); + (std::string("Starting Python backend stub: ") + log_cmd.str()) + .c_str()); pid_t pid = fork(); if (pid < 0) { @@ -974,6 +984,6 @@ StubLauncher::ShareCUDAMemoryPool( if (pb_exception.what() != std::string{""}) { throw pb_exception; } - } - #endif // TRITON_ENABLE_GPU +} +#endif // TRITON_ENABLE_GPU }}}; // namespace triton::backend::python