Log failures to /dev/termination-log

This contains simple try/catch logic to write fatal failures to /dev/termination-log

Author: joerunde

launcher/src/main.rs

@@ -13,9 +13,11 @@ use std::thread;
 use std::thread::sleep;
 use std::time::{Duration, Instant};
 use std::{fs, io};
+use std::fs::File;
 use std::env::VarError;
 use std::ffi::OsString;
 use std::os::unix::process::CommandExt;
+use String;
 use tracing::{info, warn};
 
 // In most cases this gives the best performance for inferencing
@@ -81,6 +83,20 @@ struct Args {
 }
 
 fn main() -> ExitCode {
+    // Register a panic handler up-front to write to /dev/termination-log
+    let default_hook = std::panic::take_hook();
+    std::panic::set_hook(Box::new(move |panic_info| {
+        if let Some(&s) = panic_info.payload().downcast_ref::<&str>() {
+            _ = write_termination_log(s);
+        } else if let Some(s) = panic_info.payload().downcast_ref::<String>() {
+            _ = write_termination_log(s);
+        }
+        // No else case: If we cannot get good panic info, we won't write anything to the
+        // termination log. The system logs should contain better information.
+        default_hook(panic_info);
+    }));
+
+
     // Pattern match configuration
     let args = Args::parse();
 
@@ -647,3 +663,11 @@ fn resolve_tokenizer_path(model_name: &str, revision: Option<&str>) -> Result<St
         }
     }
 }
+
+fn write_termination_log(msg: &str) -> Result<(), io::Error> {
+    // Writes a message to the termination log.
+    // Creates the logfile if it doesn't exist.
+    let mut f = File::options().write(true).create(true).open("/dev/termination-log")?;
+    writeln!(f, "{}", msg)?;
+    Ok(())
+}

router/src/main.rs

@@ -1,5 +1,8 @@
 /// Text Generation Inference external gRPC server entrypoint
 use clap::Parser;
+use std::fs::File;
+use std::io;
+use std::io::Write;
 use std::net::{IpAddr, Ipv4Addr, SocketAddr};
 use text_generation_client::ShardedClient;
 use text_generation_router::server;
@@ -48,6 +51,19 @@ struct Args {
 }
 
 fn main() -> Result<(), std::io::Error> {
+    // Register a panic handler up-front to write to /dev/termination-log
+    let default_hook = std::panic::take_hook();
+    std::panic::set_hook(Box::new(move |panic_info| {
+        if let Some(&s) = panic_info.payload().downcast_ref::<&str>() {
+            _ = write_termination_log(s);
+        } else if let Some(s) = panic_info.payload().downcast_ref::<String>() {
+            _ = write_termination_log(s);
+        }
+        // No else case: If we cannot get good panic info, we won't write anything to the
+        // termination log. The system logs should contain better information.
+        default_hook(panic_info);
+    }));
+
     // Get args
     let args = Args::parse();
 
@@ -141,3 +157,11 @@ fn main() -> Result<(), std::io::Error> {
             Ok(())
         })
 }
+
+fn write_termination_log(msg: &str) -> Result<(), io::Error> {
+    // Writes a message to the termination log.
+    // Creates the logfile if it doesn't exist.
+    let mut f = File::options().write(true).create(true).open("/dev/termination-log")?;
+    writeln!(f, "{}", msg)?;
+    Ok(())
+}

server/text_generation_server/cli.py

@@ -32,36 +32,47 @@ def serve(
     uds_path: Path = "/tmp/text-generation",
 ):
     from text_generation_server import server
+    from text_generation_server.utils.termination import write_termination_log
 
     if sharded:
-        assert (
-            os.getenv("RANK", None) is not None
-        ), "RANK must be set when sharded is True"
-        assert (
-            os.getenv("WORLD_SIZE", None) is not None
-        ), "WORLD_SIZE must be set when sharded is True"
-        assert (
-            os.getenv("MASTER_ADDR", None) is not None
-        ), "MASTER_ADDR must be set when sharded is True"
-        assert (
-            os.getenv("MASTER_PORT", None) is not None
-        ), "MASTER_PORT must be set when sharded is True"
-
-    server.serve(
-        model_name,
-        revision,
-        deployment_framework,
-        dtype,
-        # Downgrade enum into str for easier management later on
-        None if quantize is None else quantize.value,
-        max_sequence_length,
-        max_new_tokens,
-        max_batch_size,
-        batch_safety_margin,
-        sharded,
-        cuda_process_memory_fraction,
-        uds_path
-    )
+        try:
+            assert (
+                os.getenv("RANK", None) is not None
+            ), "RANK must be set when sharded is True"
+            assert (
+                os.getenv("WORLD_SIZE", None) is not None
+            ), "WORLD_SIZE must be set when sharded is True"
+            assert (
+                os.getenv("MASTER_ADDR", None) is not None
+            ), "MASTER_ADDR must be set when sharded is True"
+            assert (
+                os.getenv("MASTER_PORT", None) is not None
+            ), "MASTER_PORT must be set when sharded is True"
+        except AssertionError as e:
+            write_termination_log(str(e))
+            raise e
+
+    try:
+        server.serve(
+            model_name,
+            revision,
+            deployment_framework,
+            dtype,
+            # Downgrade enum into str for easier management later on
+            None if quantize is None else quantize.value,
+            max_sequence_length,
+            max_new_tokens,
+            max_batch_size,
+            batch_safety_margin,
+            sharded,
+            cuda_process_memory_fraction,
+            uds_path
+        )
+    except Exception as e:
+        # Any exceptions in the blocking server thread here should mean that
+        # the server terminated due to an error
+        write_termination_log(str(e))
+        raise e
 
 
 @app.command()

server/text_generation_server/utils/termination.py

@@ -0,0 +1,14 @@
+"""Utils for properly logging process termination"""
+
+
+def write_termination_log(msg: str, file: str="/dev/termination-log") -> None:
+    """Writes to the termination logfile."""
+
+    try:
+        with open(file, "w") as termination_file:
+            termination_file.write(f"{msg}\n")
+    except Exception:
+        # Ignore any errors writing to the termination logfile.
+        # Users can fall back to the stdout logs, and we don't want to pollute
+        # those with an error here.
+        pass