diff --git a/src/runtime/logging.cc b/src/runtime/logging.cc index 9c193921f93b..51dc4dc2db3d 100644 --- a/src/runtime/logging.cc +++ b/src/runtime/logging.cc @@ -18,6 +18,7 @@ */ #include +#include #include #if TVM_LOG_STACK_TRACE @@ -28,6 +29,7 @@ #include #include +#include #include #include #include @@ -95,14 +97,22 @@ int BacktraceFullCallback(void* data, uintptr_t pc, const char* filename, int li backtrace_syminfo(_bt_state, pc, BacktraceSyminfoCallback, BacktraceErrorCallback, symbol_str.get()); } - s << *symbol_str; + + if (rang::rang_implementation::isTerminal(std::cout.rdbuf())) { + // This will eventually find its way to stdout, but rang is printing to a + // regular stringstream so pipe through the isatty() result manually + rang::setControlMode(rang::control::Force); + } + s << rang::fg::yellow << *symbol_str << rang::style::reset; if (filename != nullptr) { - s << std::endl << " at " << filename; + s << std::endl << " at " << rang::fg::green << filename; + s << rang::style::reset; if (lineno != 0) { s << ":" << lineno; } } + rang::setControlMode(rang::control::Auto); // Skip tvm::backtrace and tvm::LogFatal::~LogFatal at the beginning of the trace as they don't // add anything useful to the backtrace. if (!(stack_trace->lines.size() == 0 && @@ -120,7 +130,22 @@ int BacktraceFullCallback(void* data, uintptr_t pc, const char* filename, int li std::string Backtrace() { BacktraceInfo bt; - bt.max_size = 500; + + // Limit backtrace length based on TVM_BACKTRACE_LIMIT env variable + auto user_limit_s = getenv("TVM_BACKTRACE_LIMIT"); + const auto default_limit = 500; + + if (user_limit_s == nullptr) { + bt.max_size = default_limit; + } else { + // Parse out the user-set backtrace limit + try { + bt.max_size = std::stoi(user_limit_s); + } catch (const std::invalid_argument& e) { + bt.max_size = default_limit; + } + } + if (_bt_state == nullptr) { return ""; }