Skip to content

Commit b8a5035

Browse files
yinggehyinggeh
committed
Initial impl
1 parent 87f6f2a commit b8a5035

File tree

7 files changed

+340
-236
lines changed

7 files changed

+340
-236
lines changed

‎src/message_queue.h

Lines changed: 33 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -32,7 +32,12 @@
3232
#include <boost/thread/thread_time.hpp>
3333
#include <cstddef>
3434

35+
#include "pb_exception.h"
36+
#include "pb_utils.h"
3537
#include "shm_manager.h"
38+
#ifdef TRITON_PB_STUB
39+
#include "pb_stub_log.h"
40+
#endif
3641

3742
namespace triton { namespace backend { namespace python {
3843
namespace bi = boost::interprocess;
@@ -110,7 +115,20 @@ class MessageQueue {
110115

111116
{
112117
bi::scoped_lock<bi::interprocess_mutex> lock{*MutexMutable()};
113-
Buffer()[Head()] = message;
118+
int head_idx = Head();
119+
// Additional check to avoid out of bounds read/write. Check DLIS-8378 for
120+
// additional details.
121+
if (head_idx < 0 || static_cast<size_t>(head_idx) >= Size()) {
122+
constexpr const char* error_msg =
123+
"Message queue head index out of bounds";
124+
#ifdef TRITON_PB_STUB
125+
LOG_ERROR << error_msg;
126+
#else
127+
LOG_MESSAGE(TRITONSERVER_LOG_ERROR, error_msg);
128+
#endif
129+
return;
130+
}
131+
Buffer()[head_idx] = message;
114132
HeadIncrement();
115133
}
116134
SemFullMutable()->post();
@@ -145,7 +163,20 @@ class MessageQueue {
145163
}
146164
success = true;
147165

148-
Buffer()[Head()] = message;
166+
int head_idx = Head();
167+
// Additional check to avoid out of bounds read/write. Check DLIS-8378 for
168+
// additional details.
169+
if (head_idx < 0 || static_cast<size_t>(head_idx) >= Size()) {
170+
constexpr const char* error_msg =
171+
"Message queue head index out of bounds";
172+
#ifdef TRITON_PB_STUB
173+
LOG_ERROR << error_msg;
174+
#else
175+
LOG_MESSAGE(TRITONSERVER_LOG_ERROR, error_msg);
176+
#endif
177+
return;
178+
}
179+
Buffer()[head_idx] = message;
149180
HeadIncrement();
150181
}
151182
SemFullMutable()->post();

‎src/pb_bls_cancel.cc

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -27,6 +27,7 @@
2727
#include "pb_bls_cancel.h"
2828

2929
#include "pb_stub.h"
30+
#include "pb_stub_log.h"
3031

3132
namespace triton { namespace backend { namespace python {
3233

‎src/pb_cancel.cc

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -27,6 +27,7 @@
2727
#include "pb_cancel.h"
2828

2929
#include "pb_stub.h"
30+
#include "pb_stub_log.h"
3031

3132
namespace triton { namespace backend { namespace python {
3233

‎src/pb_stub.cc

Lines changed: 1 addition & 132 deletions
Original file line numberDiff line numberDiff line change
@@ -49,6 +49,7 @@
4949
#include "pb_preferred_memory.h"
5050
#include "pb_response_iterator.h"
5151
#include "pb_string.h"
52+
#include "pb_stub_log.h"
5253
#include "pb_utils.h"
5354
#include "response_sender.h"
5455
#include "scoped_defer.h"
@@ -1569,138 +1570,6 @@ Stub::ProcessBLSResponseDecoupled(std::unique_ptr<IPCMessage>& ipc_message)
15691570
}
15701571
}
15711572

1572-
std::unique_ptr<Logger> Logger::log_instance_;
1573-
1574-
std::unique_ptr<Logger>&
1575-
Logger::GetOrCreateInstance()
1576-
{
1577-
if (Logger::log_instance_.get() == nullptr) {
1578-
Logger::log_instance_ = std::make_unique<Logger>();
1579-
}
1580-
1581-
return Logger::log_instance_;
1582-
}
1583-
1584-
// Bound function, called from the python client
1585-
void
1586-
Logger::Log(const std::string& message, LogLevel level)
1587-
{
1588-
std::unique_ptr<Stub>& stub = Stub::GetOrCreateInstance();
1589-
py::object frame = py::module_::import("inspect").attr("currentframe");
1590-
py::object caller_frame = frame();
1591-
py::object info = py::module_::import("inspect").attr("getframeinfo");
1592-
py::object caller_info = info(caller_frame);
1593-
py::object filename_python = caller_info.attr("filename");
1594-
std::string filename = filename_python.cast<std::string>();
1595-
py::object lineno = caller_info.attr("lineno");
1596-
uint32_t line = lineno.cast<uint32_t>();
1597-
1598-
if (!stub->StubToParentServiceActive()) {
1599-
Logger::GetOrCreateInstance()->Log(filename, line, level, message);
1600-
} else {
1601-
std::unique_ptr<PbLog> log_msg(new PbLog(filename, line, message, level));
1602-
stub->EnqueueLogRequest(log_msg);
1603-
}
1604-
}
1605-
1606-
// Called internally (.e.g. LOG_ERROR << "Error"; )
1607-
void
1608-
Logger::Log(
1609-
const std::string& filename, uint32_t lineno, LogLevel level,
1610-
const std::string& message)
1611-
{
1612-
// If the log monitor service is not active yet, format
1613-
// and pass messages to cerr
1614-
if (!BackendLoggingActive()) {
1615-
std::string path(filename);
1616-
size_t pos = path.rfind(std::filesystem::path::preferred_separator);
1617-
if (pos != std::string::npos) {
1618-
path = path.substr(pos + 1, std::string::npos);
1619-
}
1620-
#ifdef _WIN32
1621-
std::stringstream ss;
1622-
SYSTEMTIME system_time;
1623-
GetSystemTime(&system_time);
1624-
ss << LeadingLogChar(level) << std::setfill('0') << std::setw(2)
1625-
<< system_time.wMonth << std::setw(2) << system_time.wDay << ' '
1626-
<< std::setw(2) << system_time.wHour << ':' << std::setw(2)
1627-
<< system_time.wMinute << ':' << std::setw(2) << system_time.wSecond
1628-
<< '.' << std::setw(6) << system_time.wMilliseconds * 1000 << ' '
1629-
<< static_cast<uint32_t>(GetCurrentProcessId()) << ' ' << path << ':'
1630-
<< lineno << "] ";
1631-
#else
1632-
std::stringstream ss;
1633-
struct timeval tv;
1634-
gettimeofday(&tv, NULL);
1635-
struct tm tm_time;
1636-
gmtime_r(((time_t*)&(tv.tv_sec)), &tm_time);
1637-
ss << LeadingLogChar(level) << std::setfill('0') << std::setw(2)
1638-
<< (tm_time.tm_mon + 1) << std::setw(2) << tm_time.tm_mday << " "
1639-
<< std::setw(2) << tm_time.tm_hour << ':' << std::setw(2)
1640-
<< tm_time.tm_min << ':' << std::setw(2) << tm_time.tm_sec << "."
1641-
<< std::setw(6) << tv.tv_usec << ' ' << static_cast<uint32_t>(getpid())
1642-
<< ' ' << path << ':' << lineno << "] ";
1643-
std::cerr << ss.str() << " " << message << std::endl;
1644-
#endif
1645-
} else {
1646-
// Ensure we do not create a stub instance before it has initialized
1647-
std::unique_ptr<Stub>& stub = Stub::GetOrCreateInstance();
1648-
std::unique_ptr<PbLog> log_msg(new PbLog(filename, lineno, message, level));
1649-
stub->EnqueueLogRequest(log_msg);
1650-
}
1651-
}
1652-
1653-
void
1654-
Logger::LogInfo(const std::string& message)
1655-
{
1656-
Logger::Log(message, LogLevel::kInfo);
1657-
}
1658-
1659-
void
1660-
Logger::LogWarn(const std::string& message)
1661-
{
1662-
Logger::Log(message, LogLevel::kWarning);
1663-
}
1664-
1665-
void
1666-
Logger::LogError(const std::string& message)
1667-
{
1668-
Logger::Log(message, LogLevel::kError);
1669-
}
1670-
1671-
void
1672-
Logger::LogVerbose(const std::string& message)
1673-
{
1674-
Logger::Log(message, LogLevel::kVerbose);
1675-
}
1676-
1677-
const std::string
1678-
Logger::LeadingLogChar(const LogLevel& level)
1679-
{
1680-
switch (level) {
1681-
case LogLevel::kWarning:
1682-
return "W";
1683-
case LogLevel::kError:
1684-
return "E";
1685-
case LogLevel::kInfo:
1686-
case LogLevel::kVerbose:
1687-
default:
1688-
return "I";
1689-
}
1690-
}
1691-
1692-
void
1693-
Logger::SetBackendLoggingActive(bool status)
1694-
{
1695-
backend_logging_active_ = status;
1696-
}
1697-
1698-
bool
1699-
Logger::BackendLoggingActive()
1700-
{
1701-
return backend_logging_active_;
1702-
}
1703-
17041573
PYBIND11_EMBEDDED_MODULE(c_python_backend_utils, module)
17051574
{
17061575
py::class_<PbError, std::shared_ptr<PbError>> triton_error(

‎src/pb_stub.h

Lines changed: 0 additions & 102 deletions
Original file line numberDiff line numberDiff line change
@@ -30,8 +30,6 @@
3030
#include <pybind11/numpy.h>
3131
#include <pybind11/stl.h>
3232

33-
#include <filesystem>
34-
3533
#include "infer_request.h"
3634
#include "infer_response.h"
3735
#include "ipc_message.h"
@@ -41,7 +39,6 @@
4139
#include "pb_cancel.h"
4240
#include "pb_log.h"
4341
#include "pb_response_iterator.h"
44-
#include "pb_utils.h"
4542

4643

4744
namespace bi = boost::interprocess;
@@ -54,105 +51,6 @@ using cudaStream_t = void*;
5451

5552
namespace triton { namespace backend { namespace python {
5653

57-
#define LOG_IF_EXCEPTION(X) \
58-
do { \
59-
try { \
60-
(X); \
61-
} \
62-
catch (const PythonBackendException& pb_exception) { \
63-
LOG_INFO << pb_exception.what(); \
64-
} \
65-
} while (false)
66-
67-
#define LOG_EXCEPTION(E) \
68-
do { \
69-
LOG_INFO << E.what(); \
70-
} while (false)
71-
72-
/// Macros that use current filename and line number.
73-
#define LOG_INFO LOG_FL(__FILE__, __LINE__, LogLevel::kInfo)
74-
#define LOG_WARN LOG_FL(__FILE__, __LINE__, LogLevel::kWarning)
75-
#define LOG_ERROR LOG_FL(__FILE__, __LINE__, LogLevel::kError)
76-
#define LOG_VERBOSE LOG_FL(__FILE__, __LINE__, LogLevel::kVerbose)
77-
78-
class Logger {
79-
public:
80-
Logger() { backend_logging_active_ = false; };
81-
~Logger() { log_instance_.reset(); };
82-
/// Python client log function
83-
static void Log(const std::string& message, LogLevel level = LogLevel::kInfo);
84-
85-
/// Python client log info function
86-
static void LogInfo(const std::string& message);
87-
88-
/// Python client warning function
89-
static void LogWarn(const std::string& message);
90-
91-
/// Python client log error function
92-
static void LogError(const std::string& message);
93-
94-
/// Python client log verbose function
95-
static void LogVerbose(const std::string& message);
96-
97-
/// Internal log function
98-
void Log(
99-
const std::string& filename, uint32_t lineno, LogLevel level,
100-
const std::string& message);
101-
102-
/// Log format helper function
103-
const std::string LeadingLogChar(const LogLevel& level);
104-
105-
/// Set PYBE Logging Status
106-
void SetBackendLoggingActive(bool status);
107-
108-
/// Get PYBE Logging Status
109-
bool BackendLoggingActive();
110-
111-
/// Singleton Getter Function
112-
static std::unique_ptr<Logger>& GetOrCreateInstance();
113-
114-
DISALLOW_COPY_AND_ASSIGN(Logger);
115-
116-
/// Flush the log.
117-
void Flush() { std::cerr << std::flush; }
118-
119-
private:
120-
static std::unique_ptr<Logger> log_instance_;
121-
bool backend_logging_active_;
122-
};
123-
124-
class LogMessage {
125-
public:
126-
/// Create a log message, stripping the path down to the filename only
127-
LogMessage(const char* file, int line, LogLevel level) : level_(level)
128-
{
129-
std::string path(file);
130-
const char os_slash = std::filesystem::path::preferred_separator;
131-
size_t pos = path.rfind(os_slash);
132-
if (pos != std::string::npos) {
133-
path = path.substr(pos + 1, std::string::npos);
134-
}
135-
file_ = path;
136-
line_ = static_cast<uint32_t>(line);
137-
}
138-
/// Log message to console or send to backend (see Logger::Log for details)
139-
~LogMessage()
140-
{
141-
Logger::GetOrCreateInstance()->Log(file_, line_, level_, stream_.str());
142-
}
143-
144-
std::stringstream& stream() { return stream_; }
145-
146-
private:
147-
std::stringstream stream_;
148-
std::string file_;
149-
uint32_t line_;
150-
LogLevel level_;
151-
};
152-
153-
#define LOG_FL(FN, LN, LVL) LogMessage((char*)(FN), LN, LVL).stream()
154-
155-
15654
class ModelContext {
15755
public:
15856
// Scans and establishes path for serving the python model.

0 commit comments

Comments
 (0)