From 24bd47e88438ef2baa7e6d1e0d889c9bce81b8b4 Mon Sep 17 00:00:00 2001
From: Michal Moskal <michal@moskal.me>
Date: Tue, 23 Jan 2024 23:05:43 +0000
Subject: [PATCH] better logging

---
 aicirt/src/lib.rs      |  2 +-
 rllm/src/server/mod.rs |  9 +++++++++
 rllm/src/util.rs       | 19 +++++++++++++++++--
 3 files changed, 27 insertions(+), 3 deletions(-)

diff --git a/aicirt/src/lib.rs b/aicirt/src/lib.rs
index 69edb44b..dacefc32 100644
--- a/aicirt/src/lib.rs
+++ b/aicirt/src/lib.rs
@@ -43,7 +43,7 @@ fn daemon_format(
 
 pub fn init_log(mode: LogMode) -> Result<()> {
     let logger = match mode {
-        LogMode::Normal => Logger::try_with_env_or_str("warn")?.log_to_stdout(),
+        LogMode::Normal => Logger::try_with_env_or_str("info")?.log_to_stdout(),
         LogMode::Test => {
             Logger::try_with_env_or_str("debug")?.write_mode(WriteMode::SupportCapture)
         }
diff --git a/rllm/src/server/mod.rs b/rllm/src/server/mod.rs
index 4794f51b..2d452b61 100644
--- a/rllm/src/server/mod.rs
+++ b/rllm/src/server/mod.rs
@@ -61,6 +61,10 @@ pub struct RllmCliArgs {
     #[arg(long, short, name = "NAME=VALUE")]
     pub setting: Vec<String>,
 
+    /// Set log level (ex: "info,rllm=trace,aicirt=debug", "warn", "error")
+    #[arg(long)]
+    pub log: Option<String>,
+
     /// HuggingFace model name, URL or path starting with "./"
     #[arg(short, long, help_heading = "Model")]
     pub model: String,
@@ -468,6 +472,11 @@ fn guess_aicirt() -> Result<String> {
 
 // #[actix_web::main]
 pub async fn server_main(mut args: RllmCliArgs) -> () {
+    // we setenv, so that aicirt process also gets it
+    match &args.log {
+        Some(v) => std::env::set_var("RUST_LOG", v),
+        None => {}
+    }
     aicirt::init_log(if args.daemon {
         aicirt::LogMode::Deamon
     } else {
diff --git a/rllm/src/util.rs b/rllm/src/util.rs
index 99053e7d..1bb06a7b 100644
--- a/rllm/src/util.rs
+++ b/rllm/src/util.rs
@@ -70,12 +70,27 @@ pub fn apply_settings(settings: &Vec<String>) -> Result<()> {
     Ok(())
 }
 
+fn logging_info() -> &'static str {
+    r#"
+There are 5 logging levels: error, warn, info, debug, and trace.
+You can assign different log levels to crates:
+
+   info,rllm=trace,aicirt=debug - rllm at trace, aicirt at debug, otherwise info
+
+You can set logging levels with --log or with RUST_LOG environment variable.
+"#
+}
+
 pub fn parse_with_settings<T>() -> T
 where
     T: Parser + Args,
 {
-    let cli =
-        Command::new("CLI").after_help(format!("\n{}\n{}", all_settings(), list_tokenizers()));
+    let cli = Command::new("CLI").after_help(format!(
+        "\n{}\n{}\n{}",
+        all_settings(),
+        list_tokenizers(),
+        logging_info()
+    ));
     let cli = T::augment_args(cli);
     let matches = cli.get_matches();
     T::from_arg_matches(&matches)