Skip to content

Commit 09348d2

Browse files
committed
Implement --log-colors with always/never/auto
With auto by default Signed-off-by: Eric Curtin <[email protected]>
1 parent fb15d64 commit 09348d2

File tree

3 files changed

+89
-22
lines changed

3 files changed

+89
-22
lines changed

common/arg.cpp

Lines changed: 44 additions & 22 deletions
Original file line numberDiff line numberDiff line change
@@ -1263,6 +1263,18 @@ static std::string list_builtin_chat_templates() {
12631263
return msg.str();
12641264
}
12651265

1266+
static bool is_truthy(const std::string & value) {
1267+
return value == "always" || value == "on" || value == "enabled" || value == "1";
1268+
}
1269+
1270+
static bool is_falsey(const std::string & value) {
1271+
return value == "never" || value == "off" || value == "disabled" || value == "0";
1272+
}
1273+
1274+
static bool is_autoy(const std::string & value) {
1275+
return value == "auto" || value == "-1";
1276+
}
1277+
12661278
common_params_context common_params_parser_init(common_params & params, llama_example ex, void(*print_usage)(int, char **)) {
12671279
// load dynamic backends
12681280
ggml_backend_load_all();
@@ -1544,21 +1556,22 @@ common_params_context common_params_parser_init(common_params & params, llama_ex
15441556
params.n_chunks = value;
15451557
}
15461558
).set_examples({LLAMA_EXAMPLE_IMATRIX, LLAMA_EXAMPLE_PERPLEXITY, LLAMA_EXAMPLE_RETRIEVAL}));
1547-
add_opt(common_arg(
1548-
{"-fa", "--flash-attn"}, "FA",
1549-
string_format("set Flash Attention use ('on', 'off', or 'auto', default: '%s')", llama_flash_attn_type_name(params.flash_attn_type)),
1550-
[](common_params & params, const std::string & value) {
1551-
if (value == "on" || value == "enabled" || value == "1") {
1552-
params.flash_attn_type = LLAMA_FLASH_ATTN_TYPE_ENABLED;
1553-
} else if (value == "off" || value == "disabled" || value == "0") {
1554-
params.flash_attn_type = LLAMA_FLASH_ATTN_TYPE_DISABLED;
1555-
} else if (value == "auto" || value == "-1") {
1556-
params.flash_attn_type = LLAMA_FLASH_ATTN_TYPE_AUTO;
1557-
} else {
1558-
throw std::runtime_error(string_format("error: unkown value for --flash-attn: '%s'\n", value.c_str()));
1559-
}
1560-
}
1561-
).set_env("LLAMA_ARG_FLASH_ATTN"));
1559+
add_opt(common_arg({ "-fa", "--flash-attn" }, "FA",
1560+
string_format("set Flash Attention use ('on', 'off', or 'auto', default: '%s')",
1561+
llama_flash_attn_type_name(params.flash_attn_type)),
1562+
[](common_params & params, const std::string & value) {
1563+
if (is_truthy(value)) {
1564+
params.flash_attn_type = LLAMA_FLASH_ATTN_TYPE_ENABLED;
1565+
} else if (is_falsey(value)) {
1566+
params.flash_attn_type = LLAMA_FLASH_ATTN_TYPE_DISABLED;
1567+
} else if (is_autoy(value)) {
1568+
params.flash_attn_type = LLAMA_FLASH_ATTN_TYPE_AUTO;
1569+
} else {
1570+
throw std::runtime_error(
1571+
string_format("error: unkown value for --flash-attn: '%s'\n", value.c_str()));
1572+
}
1573+
})
1574+
.set_env("LLAMA_ARG_FLASH_ATTN"));
15621575
add_opt(common_arg(
15631576
{"-p", "--prompt"}, "PROMPT",
15641577
"prompt to start generation with; for system message, use -sys",
@@ -3134,13 +3147,22 @@ common_params_context common_params_parser_init(common_params & params, llama_ex
31343147
common_log_set_file(common_log_main(), value.c_str());
31353148
}
31363149
));
3137-
add_opt(common_arg(
3138-
{"--log-colors"},
3139-
"Enable colored logging",
3140-
[](common_params &) {
3141-
common_log_set_colors(common_log_main(), true);
3142-
}
3143-
).set_env("LLAMA_LOG_COLORS"));
3150+
add_opt(common_arg({ "--log-colors" }, "[always|never|auto]",
3151+
"Enable colored logging ('on', 'off', or 'auto', default: 'auto')\n"
3152+
"'auto' enables colors when output is to a terminal",
3153+
[](common_params &, const std::string & value) {
3154+
if (is_truthy(value)) {
3155+
common_log_set_colors(common_log_main(), true);
3156+
} else if (is_falsey(value)) {
3157+
common_log_set_colors(common_log_main(), false);
3158+
} else if (is_autoy(value)) {
3159+
common_log_set_colors_auto(common_log_main());
3160+
} else {
3161+
throw std::invalid_argument(
3162+
string_format("error: unkown value for --log-colors: '%s'\n", value.c_str()));
3163+
}
3164+
})
3165+
.set_env("LLAMA_LOG_COLORS"));
31443166
add_opt(common_arg(
31453167
{"-v", "--verbose", "--log-verbose"},
31463168
"Set verbosity level to infinity (i.e. log all messages, useful for debugging)",

common/log.cpp

Lines changed: 44 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -4,17 +4,52 @@
44
#include <condition_variable>
55
#include <cstdarg>
66
#include <cstdio>
7+
#include <cstdlib>
8+
#include <cstring>
79
#include <mutex>
810
#include <sstream>
911
#include <thread>
1012
#include <vector>
1113

14+
#if defined(_WIN32)
15+
# include <io.h>
16+
# include <windows.h>
17+
# define isatty _isatty
18+
# define fileno _fileno
19+
#else
20+
# include <unistd.h>
21+
#endif
22+
1223
int common_log_verbosity_thold = LOG_DEFAULT_LLAMA;
1324

1425
void common_log_set_verbosity_thold(int verbosity) {
1526
common_log_verbosity_thold = verbosity;
1627
}
1728

29+
// Auto-detect if colors should be enabled based on terminal and environment
30+
static bool common_log_should_use_colors_auto() {
31+
// Check NO_COLOR environment variable (https://no-color.org/)
32+
if (const char * no_color = std::getenv("NO_COLOR")) {
33+
if (no_color[0] != '\0') {
34+
return false;
35+
}
36+
}
37+
38+
// Check TERM environment variable
39+
if (const char * term = std::getenv("TERM")) {
40+
if (std::strcmp(term, "dumb") == 0) {
41+
return false;
42+
}
43+
}
44+
45+
// Check if stdout and stderr are connected to a terminal
46+
// We check both because log messages can go to either
47+
bool stdout_is_tty = isatty(fileno(stdout)) != 0;
48+
bool stderr_is_tty = isatty(fileno(stderr)) != 0;
49+
50+
return stdout_is_tty || stderr_is_tty;
51+
}
52+
1853
static int64_t t_us() {
1954
return std::chrono::duration_cast<std::chrono::microseconds>(std::chrono::system_clock::now().time_since_epoch()).count();
2055
}
@@ -353,6 +388,11 @@ struct common_log * common_log_init() {
353388

354389
struct common_log * common_log_main() {
355390
static struct common_log log;
391+
static std::once_flag init_flag;
392+
std::call_once(init_flag, [&]() {
393+
// Set default to auto-detect colors
394+
log.set_colors(common_log_should_use_colors_auto());
395+
});
356396

357397
return &log;
358398
}
@@ -384,6 +424,10 @@ void common_log_set_colors(struct common_log * log, bool colors) {
384424
log->set_colors(colors);
385425
}
386426

427+
void common_log_set_colors_auto(struct common_log * log) {
428+
log->set_colors(common_log_should_use_colors_auto());
429+
}
430+
387431
void common_log_set_prefix(struct common_log * log, bool prefix) {
388432
log->set_prefix(prefix);
389433
}

common/log.h

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -67,6 +67,7 @@ void common_log_add(struct common_log * log, enum ggml_log_level level, const ch
6767

6868
void common_log_set_file (struct common_log * log, const char * file); // not thread-safe
6969
void common_log_set_colors (struct common_log * log, bool colors); // not thread-safe
70+
void common_log_set_colors_auto(struct common_log * log); // auto-detect colors based on terminal, not thread-safe
7071
void common_log_set_prefix (struct common_log * log, bool prefix); // whether to output prefix to each log
7172
void common_log_set_timestamps(struct common_log * log, bool timestamps); // whether to output timestamps in the prefix
7273

0 commit comments

Comments
 (0)