Adds flag --history-only to "hq journal stream"

Author: spirali

crates/hyperqueue/src/client/commands/journal/mod.rs

@@ -6,7 +6,7 @@ use crate::common::utils::str::pluralize;
 use crate::rpc_call;
 use crate::server::bootstrap::get_client_session;
 use crate::server::event::journal::JournalReader;
-use crate::transfer::messages::{FromClientMessage, ToClientMessage};
+use crate::transfer::messages::{FromClientMessage, StreamEvents, ToClientMessage};
 use anyhow::anyhow;
 use clap::{Parser, ValueHint};
 use std::io::{BufWriter, Write};
@@ -25,8 +25,8 @@ enum JournalCommand {
     /// Events will be exported to `stdout`, you can redirect it e.g. to a file.
     Export(ExportOpts),
 
-    /// Live stream events from the server.
-    Stream,
+    /// Stream events from a running server.
+    Stream(StreamOpts),
 
     /// Connect to a server and remove completed tasks and non-active workers from journal
     Prune,
@@ -43,20 +43,30 @@ struct ExportOpts {
     journal: PathBuf,
 }
 
+#[derive(Parser)]
+struct StreamOpts {
+    /// If enabled, server terminates the connection when all currents events
+    /// are sent.
+    #[arg(long)]
+    history_only: bool,
+}
+
 pub async fn command_journal(gsettings: &GlobalSettings, opts: JournalOpts) -> anyhow::Result<()> {
     match opts.command {
         JournalCommand::Export(opts) => export_json(opts),
-        JournalCommand::Stream => stream_json(gsettings).await,
+        JournalCommand::Stream(opts) => stream_json(gsettings, opts).await,
         JournalCommand::Prune => prune_journal(gsettings).await,
         JournalCommand::Flush => flush_journal(gsettings).await,
     }
 }
 
-async fn stream_json(gsettings: &GlobalSettings) -> anyhow::Result<()> {
+async fn stream_json(gsettings: &GlobalSettings, opts: StreamOpts) -> anyhow::Result<()> {
     let mut connection = get_client_session(gsettings.server_directory()).await?;
     connection
         .connection()
-        .send(FromClientMessage::StreamEvents)
+        .send(FromClientMessage::StreamEvents(StreamEvents {
+            history_only: opts.history_only,
+        }))
         .await?;
     let stdout = std::io::stdout();
     let stdout = stdout.lock();

crates/hyperqueue/src/dashboard/data/fetch.rs

@@ -1,6 +1,6 @@
 use crate::server::event::Event;
 use crate::transfer::connection::ClientSession;
-use crate::transfer::messages::{FromClientMessage, ToClientMessage};
+use crate::transfer::messages::{FromClientMessage, StreamEvents, ToClientMessage};
 use std::time::Duration;
 use tokio::sync::mpsc::Sender;
 
@@ -10,7 +10,9 @@ pub async fn create_data_fetch_process(
 ) -> anyhow::Result<()> {
     session
         .connection()
-        .send(FromClientMessage::StreamEvents)
+        .send(FromClientMessage::StreamEvents(StreamEvents {
+            history_only: false,
+        }))
         .await?;
 
     const CAPACITY: usize = 1024;

crates/hyperqueue/src/server/client/mod.rs

@@ -75,21 +75,26 @@ async fn handle_client(
     Ok(())
 }
 
-async fn stream_events<
-    Tx: Sink<ToClientMessage, Error = HqError> + Unpin + 'static,
-    Rx: Stream<Item = crate::Result<FromClientMessage>> + Unpin,
->(
+async fn stream_history_events<Tx: Sink<ToClientMessage, Error = HqError> + Unpin + 'static>(
     tx: &mut Tx,
-    rx: &mut Rx,
     mut history: mpsc::UnboundedReceiver<Event>,
-    mut current: mpsc::UnboundedReceiver<Event>,
 ) {
     log::debug!("Resending history started");
     while let Some(e) = history.recv().await {
         if tx.send(ToClientMessage::Event(e)).await.is_err() {
             return;
         }
     }
+}
+
+async fn stream_events<
+    Tx: Sink<ToClientMessage, Error = HqError> + Unpin + 'static,
+    Rx: Stream<Item = crate::Result<FromClientMessage>> + Unpin,
+>(
+    tx: &mut Tx,
+    rx: &mut Rx,
+    mut current: mpsc::UnboundedReceiver<Event>,
+) {
     log::debug!("History streaming completed");
     loop {
         let r = tokio::select! {
@@ -164,18 +169,25 @@ pub async fn client_rpc_loop<
                     FromClientMessage::CloseJob(msg) => {
                         handle_job_close(&state_ref, senders, &msg.selector).await
                     }
-                    FromClientMessage::StreamEvents => {
+                    FromClientMessage::StreamEvents(msg) => {
                         log::debug!("Start streaming events to client");
                         /* We create two event queues, one for historic events and one for live events
                         So while historic events are loaded from the file and streamed, live events are already
                         collected and sent immediately the historic events are sent */
                         let (tx1, rx1) = mpsc::unbounded_channel::<Event>();
-                        let (tx2, rx2) = mpsc::unbounded_channel::<Event>();
-                        let listener_id = senders.events.register_listener(tx1, tx2);
-
-                        stream_events(&mut tx, &mut rx, rx1, rx2).await;
-
-                        senders.events.unregister_listener(listener_id);
+                        let live = if !msg.history_only {
+                            let (tx2, rx2) = mpsc::unbounded_channel::<Event>();
+                            let listener_id = senders.events.register_listener(tx2);
+                            Some((rx2, listener_id))
+                        } else {
+                            None
+                        };
+                        senders.events.replay_journal(tx1);
+                        stream_history_events(&mut tx, rx1).await;
+                        if let Some((rx2, listener_id)) = live {
+                            stream_events(&mut tx, &mut rx, rx2).await;
+                            senders.events.unregister_listener(listener_id);
+                        }
                         break;
                     }
                     FromClientMessage::ServerInfo => {

crates/hyperqueue/src/server/event/journal/stream.rs

@@ -15,7 +15,7 @@ use tokio::sync::mpsc;
 
 pub enum EventStreamMessage {
     Event(Event),
-    RegisterListener(mpsc::UnboundedSender<Event>),
+    ReplayJournal(mpsc::UnboundedSender<Event>),
     PruneJournal {
         callback: tokio::sync::oneshot::Sender<()>,
         live_jobs: Set<JobId>,
@@ -88,7 +88,7 @@ async fn streaming_process(
                             break
                         }
                     }
-                    Some(EventStreamMessage::RegisterListener(tx)) => {
+                    Some(EventStreamMessage::ReplayJournal(tx)) => {
                         /* We are blocking the thread here, but it is intended.
                            But we are blocking just a thread managing log file, not the whole HQ
                            And while this read is performed, we cannot allow modification of the file,

crates/hyperqueue/src/server/event/streamer.rs

@@ -169,11 +169,19 @@ impl EventStreamer {
         });
     }
 
-    pub fn register_listener(
-        &self,
-        history_sender: mpsc::UnboundedSender<Event>,
-        current_sender: mpsc::UnboundedSender<Event>,
-    ) -> u32 {
+    pub fn replay_journal(&self, history_sender: mpsc::UnboundedSender<Event>) {
+        let inner = self.inner.get();
+        if let Some(ref streamer) = inner.storage_sender {
+            if streamer
+                .send(EventStreamMessage::ReplayJournal(history_sender))
+                .is_err()
+            {
+                log::error!("Event streaming queue has been closed.");
+            }
+        }
+    }
+
+    pub fn register_listener(&self, current_sender: mpsc::UnboundedSender<Event>) -> u32 {
         let mut inner = self.inner.get_mut();
         let listener_id = inner
             .client_listeners
@@ -183,14 +191,6 @@ impl EventStreamer {
             .unwrap_or(0)
             + 1;
         inner.client_listeners.push((current_sender, listener_id));
-        if let Some(ref streamer) = inner.storage_sender {
-            if streamer
-                .send(EventStreamMessage::RegisterListener(history_sender))
-                .is_err()
-            {
-                log::error!("Event streaming queue has been closed.");
-            }
-        }
         listener_id
     }
 

crates/hyperqueue/src/transfer/messages.rs

@@ -41,7 +41,7 @@ pub enum FromClientMessage {
     // This command switches the connection into streaming connection,
     // it will no longer reacts to any other client messages
     // and client will only receive ToClientMessage::Event
-    StreamEvents,
+    StreamEvents(StreamEvents),
     PruneJournal,
     FlushJournal,
 }
@@ -86,6 +86,11 @@ pub struct TaskKindProgram {
     pub task_dir: bool,
 }
 
+#[derive(Serialize, Deserialize, Debug, Clone)]
+pub struct StreamEvents {
+    pub history_only: bool,
+}
+
 #[derive(Serialize, Deserialize, Debug, Clone)]
 pub enum TaskKind {
     ExternalProgram(TaskKindProgram),

tests/autoalloc/test_autoalloc.py

@@ -88,15 +88,17 @@ def test_slurm_queue_sbatch_additional_output(hq_env: HqEnv):
     class Handler(SlurmCommandHandler):
         async def handle_submit(self, input: CommandInput) -> CommandOutput:
             output = await super().handle_submit(input)
-            return response(f"""
+            return response(
+                f"""
 No reservation for this job
 --> Verifying valid submit host (login)...OK
 --> Verifying valid jobname...OK
 --> Verifying valid ssh keys...OK
 --> Verifying access to desired queue (normal)...OK
 --> Checking available allocation...OK
 {output.stdout}
-""")
+"""
+            )
 
     with MockJobManager(hq_env, Handler(DefaultManager())):
         hq_env.start_server()

tests/test_events.py

@@ -64,6 +64,14 @@ def test_worker_stream_events2(hq_env: HqEnv, tmp_path):
         assert events[2]["event"]["desc"]["name"] == "uname"
 
 
+def test_worker_stream_history_only(hq_env: HqEnv, tmp_path):
+    journal = tmp_path.joinpath("test.journal")
+    hq_env.start_server(args=["--journal", journal])
+    r = hq_env.command(["journal", "stream", "--history-only"])
+    msg = json.loads(r)
+    assert msg["event"]["type"] == "server-start"
+
+
 def test_worker_connected_event(hq_env: HqEnv):
     def body():
         hq_env.start_worker()
@@ -154,21 +162,21 @@ def body():
         with hq_env.mock.mock_program_with_code(
             "rocm-smi",
             """
-import json
-data = {
-    "card0": {
-        "GPU use (%)": "1.5",
-        "GPU memory use (%)": "12.5",
-        "PCI Bus": "FOOBAR1"
-    },
-    "card1": {
-        "GPU use (%)": "12.5",
-        "GPU memory use (%)": "64.0",
-        "PCI Bus": "FOOBAR2"
+    import json
+    data = {
+        "card0": {
+            "GPU use (%)": "1.5",
+            "GPU memory use (%)": "12.5",
+            "PCI Bus": "FOOBAR1"
+        },
+        "card1": {
+            "GPU use (%)": "12.5",
+            "GPU memory use (%)": "64.0",
+            "PCI Bus": "FOOBAR2"
+        }
     }
-}
-print(json.dumps(data))
-        """,
+    print(json.dumps(data))
+            """,
         ):
             hq_env.start_worker(args=["--overview-interval", "10ms", "--resource", "gpus/amd=[0]"])
             wait_for_worker_state(hq_env, 1, "RUNNING")