Split Actor logs in their own files

Summary:
- Each actor type will create a new file suffixed with the actor name
- Each log will contain the proc and rank id of the actor. This will simplify filtering to logs of a single actor
- Logs not coming from an actor will be dumped in a common file suffixed 'monarch'

TODO:
- Support Python Actors
- Support multiple configurable streams for Local

Differential Revision: D78101221

Author: pablorfb-meta

hyperactor/src/proc.rs

@@ -919,7 +919,15 @@ impl<A: Actor> Instance<A> {
     }
 
     async fn serve(mut self, mut actor: A) {
-        let result = self.run_actor_tree(&mut actor).await;
+        let actor_id = self.cell.actor_id().clone();
+        let filename = actor_id.name().to_string();
+        let prefix = format!("{}_actor[{}]", actor_id.proc_id(), actor_id.rank());
+
+        let result =
+            hyperactor_telemetry::with_logging_file_and_prefix(&filename, &prefix, || async {
+                self.run_actor_tree(&mut actor).await
+            })
+            .await;
 
         let actor_status = match result {
             Ok(_) => ActorStatus::Stopped,
@@ -1026,7 +1034,7 @@ impl<A: Actor> Instance<A> {
     /// Initialize and run the actor until it fails or is stopped.
     async fn run(&mut self, actor: &mut A) -> Result<(), ActorError> {
         hyperactor_telemetry::declare_static_counter!(MESSAGES_RECEIVED, "actor.messages_received");
-        tracing::debug!("entering actor loop: {}", self.self_id());
+        tracing::debug!("entering actor loop for {}", self.cell.actor_id());
 
         self.change_status(ActorStatus::Initializing);
         actor

hyperactor_telemetry/src/lib.rs

@@ -36,6 +36,7 @@ mod otel;
 mod pool;
 pub mod recorder;
 mod spool;
+use std::cell::RefCell;
 use std::io::IsTerminal;
 use std::io::Write;
 use std::str::FromStr;
@@ -61,6 +62,11 @@ use tracing_subscriber::Layer;
 use tracing_subscriber::filter::LevelFilter;
 use tracing_subscriber::filter::Targets;
 use tracing_subscriber::fmt;
+use tracing_subscriber::fmt::FormatEvent;
+use tracing_subscriber::fmt::FormatFields;
+use tracing_subscriber::fmt::MakeWriter;
+use tracing_subscriber::fmt::format::Writer;
+use tracing_subscriber::registry::LookupSpan;
 
 use crate::recorder::Recorder;
 
@@ -81,26 +87,28 @@ impl TelemetryClock for DefaultTelemetryClock {
     }
 }
 
-// Need to keep this around so that the tracing subscriber doesn't drop the writer.
-lazy_static! {
-    static ref WRITER_GUARD: Arc<(NonBlocking, WorkerGuard)> = {
-        let writer: Box<dyn Write + Send> = match env::Env::current() {
-            env::Env::Local | env::Env::Test | env::Env::MastEmulator => {
+fn writer(filename: &str) -> Box<dyn Write + Send> {
+    match env::Env::current() {
+        env::Env::Test | env::Env::MastEmulator => Box::new(std::io::stderr()),
+        env::Env::Local | env::Env::Mast => match RollingFileAppender::builder()
+            .rotation(Rotation::HOURLY)
+            .filename_prefix(format!("dedicated_log_{}", filename))
+            .filename_suffix("log")
+            .build("/logs/")
+        {
+            Ok(file) => Box::new(file),
+            Err(e) => {
+                tracing::warn!("unable to create custom log file: {}", e);
                 Box::new(std::io::stderr())
             }
-            env::Env::Mast => match RollingFileAppender::builder()
-                .rotation(Rotation::HOURLY)
-                .filename_prefix("dedicated_log_monarch")
-                .filename_suffix("log")
-                .build("/logs/")
-            {
-                Ok(file) => Box::new(file),
-                Err(e) => {
-                    tracing::warn!("unable to create custom log file: {}", e);
-                    Box::new(std::io::stderr())
-                }
-            },
-        };
+        },
+    }
+}
+
+// Need to keep this around so that the tracing subscriber doesn't drop the writer.
+lazy_static! {
+    static ref DEFAULT_WRITER_GUARD: Arc<(NonBlocking, WorkerGuard)> = {
+        let writer: Box<dyn Write + Send> = writer("monarch");
         return Arc::new(
             tracing_appender::non_blocking::NonBlockingBuilder::default()
                 .lossy(false)
@@ -415,6 +423,159 @@ macro_rules! declare_static_histogram {
     };
 }
 
+tokio::task_local! {
+    static TASK_LOG_FILENAME: RefCell<Option<Box<dyn Write + Send>>>;
+    static TASK_LOG_PREFIX: RefCell<Option<String>>;
+}
+
+/// Set both logging file and prefix for the current tokio task using a scope.
+/// This creates a scope where all log messages will be written to the specified file with the given prefix.
+pub async fn with_logging_file_and_prefix<F, Fut>(filename: &str, prefix: &str, f: F) -> Fut::Output
+where
+    F: FnOnce() -> Fut,
+    Fut: std::future::Future,
+{
+    let writer_cell = RefCell::new(Some(writer(filename)));
+    let prefix_cell = RefCell::new(Some(prefix.to_string()));
+
+    TASK_LOG_FILENAME
+        .scope(writer_cell, async move {
+            TASK_LOG_PREFIX
+                .scope(prefix_cell, async move { f().await })
+                .await
+        })
+        .await
+}
+
+/// A custom formatter that prepends task-local prefixes to log messages.
+/// Wraps the Glog formatter and adds task-specific prefix.
+pub struct PrefixedGlogFormatter {
+    inner: Glog<LocalTime>,
+}
+
+impl PrefixedGlogFormatter {
+    pub fn new() -> Self {
+        Self {
+            inner: Glog::default().with_timer(LocalTime::default()),
+        }
+    }
+}
+
+impl Default for PrefixedGlogFormatter {
+    fn default() -> Self {
+        Self::new()
+    }
+}
+
+impl<S, N> FormatEvent<S, N> for PrefixedGlogFormatter
+where
+    S: tracing::Subscriber + for<'a> LookupSpan<'a>,
+    N: for<'a> FormatFields<'a> + 'static,
+{
+    fn format_event(
+        &self,
+        ctx: &tracing_subscriber::fmt::FmtContext<'_, S, N>,
+        mut writer: Writer<'_>,
+        event: &tracing::Event<'_>,
+    ) -> std::fmt::Result {
+        let prefix = TASK_LOG_PREFIX
+            .try_with(|log_prefix| log_prefix.borrow().clone())
+            .unwrap_or(None);
+        if let Some(prefix) = prefix {
+            write!(writer, "[{}]", prefix)?;
+        }
+
+        self.inner.format_event(ctx, writer, event)
+    }
+}
+
+/// A custom writer that routes log events to task-specific files with fallback.
+/// Checks if the current task has a specific log file configured,
+/// and if not, falls back to the provided fallback writer.
+#[derive(Debug, Clone)]
+pub struct TaskRoutingWriter<W> {
+    fallback_writer: W,
+}
+
+impl<W> TaskRoutingWriter<W> {
+    /// Create a new TaskRoutingWriter with the specified fallback writer.
+    pub fn new(fallback_writer: W) -> Self {
+        Self { fallback_writer }
+    }
+}
+
+/// Wrapper that implements Write and routes to custom sink if defined, otherwise to fallback.
+pub struct CustomWriter<W> {
+    sink: Sink<W>,
+}
+
+enum Sink<W> {
+    Custom,
+    Fallback(W),
+}
+
+impl<W: std::io::Write> std::io::Write for CustomWriter<W> {
+    fn write(&mut self, buf: &[u8]) -> std::io::Result<usize> {
+        match &mut self.sink {
+            Sink::Custom => TASK_LOG_FILENAME
+                .try_with(|writer_ref| {
+                    let mut writer_borrow = writer_ref.borrow_mut();
+                    if let Some(ref mut writer) = *writer_borrow {
+                        writer.write(buf)
+                    } else {
+                        Err(std::io::Error::new(
+                            std::io::ErrorKind::NotFound,
+                            "Expected custom taks writer but found none.",
+                        ))
+                    }
+                })
+                .unwrap_or_else(|_| {
+                    Err(std::io::Error::new(
+                        std::io::ErrorKind::NotFound,
+                        "Unable to access task writer.",
+                    ))
+                }),
+            Sink::Fallback(writer) => writer.write(buf),
+        }
+    }
+
+    fn flush(&mut self) -> std::io::Result<()> {
+        match &mut self.sink {
+            Sink::Custom => TASK_LOG_FILENAME
+                .try_with(|log| {
+                    let mut log_borrow = log.borrow_mut();
+                    if let Some(ref mut writer) = *log_borrow {
+                        writer.flush()
+                    } else {
+                        Ok(())
+                    }
+                })
+                .unwrap_or(Ok(())),
+            Sink::Fallback(writer) => writer.flush(),
+        }
+    }
+}
+
+impl<'a, W> MakeWriter<'a> for TaskRoutingWriter<W>
+where
+    W: for<'writer> MakeWriter<'writer> + Clone,
+{
+    type Writer = CustomWriter<<W as MakeWriter<'a>>::Writer>;
+
+    fn make_writer(&'a self) -> Self::Writer {
+        match TASK_LOG_FILENAME.try_with(|log| log.borrow().is_some()) {
+            Ok(_) => CustomWriter { sink: Sink::Custom },
+            Err(_) => CustomWriter {
+                sink: Sink::Fallback(self.fallback_writer.make_writer()),
+            },
+        }
+    }
+
+    fn make_writer_for(&'a self, _: &tracing::Metadata<'_>) -> Self::Writer {
+        self.make_writer()
+    }
+}
+
 /// Set up logging based on the given execution environment. We specialize logging based on how the
 /// logs are consumed. The destination scuba table is specialized based on the execution environment.
 /// mast -> monarch_tracing/prod
@@ -432,10 +593,9 @@ pub fn initialize_logging(clock: impl TelemetryClock + Send + 'static) {
         env::Env::Test => "debug",
     };
 
-    let writer: &NonBlocking = &WRITER_GUARD.0;
     let glog = fmt::Layer::default()
-        .with_writer(writer.clone())
-        .event_format(Glog::default().with_timer(LocalTime::default()))
+        .with_writer(TaskRoutingWriter::new(DEFAULT_WRITER_GUARD.0.clone()))
+        .event_format(PrefixedGlogFormatter::new())
         .fmt_fields(GlogFields::default().compact())
         .with_ansi(std::io::stderr().is_terminal())
         .with_filter(