Implement hq server wait command

Author: Kobzol

CHANGELOG.md

@@ -20,6 +20,8 @@
 * Resource policy `compact!` is now allowed to take fractional resource request.
 * There is a new command `hq alloc cat <alloc-id> <stdout/stderr>`, which can be used
   to debug the output of allocations submitted by the automatic allocator.
+* There is a new command `hq server wait` that repeatedly tries to connect to a server with a configurable timeout.
+  This is useful for deployment scripts that need to wait for server availability.
 
 ### Fixes
 

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

@@ -13,6 +13,7 @@ use crate::server::bootstrap::{
 use crate::transfer::auth::generate_key;
 use crate::transfer::messages::{FromClientMessage, ToClientMessage};
 use clap::Parser;
+use humantime::format_duration;
 use std::path::PathBuf;
 use std::sync::Arc;
 use std::time::Duration;
@@ -71,6 +72,8 @@ pub enum ServerCommand {
     GenerateAccess(GenerateAccessOpts),
     /// Dump internal scheduler info into a file
     DebugDump(DebugDumpOpts),
+    /// Wait until the server becomes available
+    Wait(WaitOpts),
 }
 
 #[derive(Parser)]
@@ -141,13 +144,28 @@ pub struct DebugDumpOpts {
     path: PathBuf,
 }
 
+#[derive(Parser)]
+pub struct WaitOpts {
+    /// Timeout after which to stop trying to connect.
+    ///
+    /// If we cannot connect until the timeout expires, the wait command will fail.
+    #[arg(
+        long,
+        value_parser = parse_hms_or_human_time,
+        default_value = "5m",
+        help = duration_doc!("Timeout after which to stop trying to connect.")
+    )]
+    timeout: Duration,
+}
+
 pub async fn command_server(gsettings: &GlobalSettings, opts: ServerOpts) -> anyhow::Result<()> {
     match opts.subcmd {
         ServerCommand::Start(opts) => start_server(gsettings, opts).await,
         ServerCommand::Stop(opts) => stop_server(gsettings, opts).await,
         ServerCommand::Info(opts) => command_server_info(gsettings, opts).await,
         ServerCommand::GenerateAccess(opts) => command_server_generate_access(gsettings, opts),
         ServerCommand::DebugDump(opts) => debug_dump(gsettings, opts).await,
+        ServerCommand::Wait(opts) => wait_for_server(gsettings, opts).await,
     }
 }
 
@@ -282,3 +300,40 @@ fn command_server_generate_access(
     }
     Ok(())
 }
+
+async fn wait_for_server(gsettings: &GlobalSettings, opts: WaitOpts) -> anyhow::Result<()> {
+    let retry_interval = Duration::from_secs(5);
+
+    log::info!(
+        "Waiting for server to become available (timeout: {})...",
+        format_duration(opts.timeout)
+    );
+
+    let result = tokio::time::timeout(opts.timeout, async {
+        loop {
+            match get_client_session(gsettings.server_directory()).await {
+                Ok(_session) => {
+                    log::info!("Successfully connected to server");
+                    break;
+                }
+                Err(_) => {
+                    log::debug!(
+                        "Server not yet available, retrying in {}s...",
+                        retry_interval.as_secs()
+                    );
+                    tokio::time::sleep(retry_interval).await;
+                }
+            }
+        }
+    })
+    .await;
+
+    match result {
+        Ok(()) => Ok(()),
+        Err(_) => Err(anyhow::anyhow!(
+            "Timeout after {}: Could not connect to server at `{}`. Make sure that the server is running.",
+            format_duration(opts.timeout),
+            gsettings.server_directory().display()
+        )),
+    }
+}

docs/deployment/server.md

@@ -127,6 +127,28 @@ The [`hq journal prune`](cli:hq.journal.prune) command removes all completed job
 
 The [`hq journal flush`](cli:hq.journal.flush) command will force the server to flush the journal, so that the latest state of affairs is persisted to disk. It is mainly useful for testing or if you are going to run `hq journal export` while a server is running (however, it is usually better to use `hq journal stream`).
 
+## Waiting for server availability
+
+If you need to wait for a server to become available (for example when coordinating server startup in scripts), 
+you can use the [`hq server wait`](cli:hq.server.wait) command:
+
+```bash
+$ hq server wait
+```
+
+This command will repeatedly attempt to connect to the server (every 5 seconds) until it succeeds or until a timeout 
+is reached. By default, it will wait for up to 5 minutes, but you can specify a custom timeout[^1]:
+
+```bash
+# Wait for up to 2 minutes
+$ hq server wait --timeout 2m
+```
+
+[^1]: You can use various [shortcuts](../cli/shortcuts.md#duration) for the timeout duration.
+
+This is particularly useful in deployment scripts where you start a server and then need to ensure it's ready before 
+proceeding with other operations like connecting workers or submitting jobs.
+
 ## Stopping the server
 
 You can stop a running server with the [`hq server stop`](cli:hq.server.stop) command:

tests/test_server.py

@@ -4,6 +4,7 @@
 import signal
 import socket
 import subprocess
+import time
 from typing import List
 
 import pytest
@@ -258,3 +259,52 @@ def test_server_debug_dump(hq_env: HqEnv):
     with open("out.json") as f:
         result = json.loads(f.read())
     assert "tako" in result
+
+
+def test_server_wait_timeout_no_server(hq_env: HqEnv, tmp_path):
+    """Test that 'hq server wait' times out when no server is running"""
+    start_time = time.time()
+
+    with pytest.raises(Exception, match=r"Timeout after 2s.*Could not connect to server"):
+        hq_env.command(["server", "wait", "--timeout", "2s", "--server-dir", str(tmp_path)], use_server_dir=False)
+
+    elapsed = time.time() - start_time
+
+    # Should timeout after roughly 2 seconds (allow some tolerance)
+    assert 1.5 < elapsed < 4.0
+
+
+def test_server_wait_success_immediate(hq_env: HqEnv):
+    """Test that 'hq server wait' succeeds immediately when server is already running"""
+    hq_env.start_server()
+
+    start_time = time.time()
+    hq_env.command(["server", "wait", "--timeout", "5s"])
+    elapsed = time.time() - start_time
+
+    # Should succeed quickly (less than 5s timeout)
+    assert elapsed < 4.0
+
+
+def test_server_wait_success_delayed(hq_env: HqEnv, tmp_path):
+    """Test that 'hq server wait' succeeds when server starts during the wait period"""
+
+    def delayed_server_start():
+        time.sleep(2)
+        hq_env.start_server(server_dir=str(tmp_path))
+
+    import threading
+
+    server_thread = threading.Thread(target=delayed_server_start)
+
+    start_time = time.time()
+    server_thread.start()
+
+    # Should succeed when server starts after ~2s
+    hq_env.command(["server", "wait", "--timeout", "10s", "--server-dir", str(tmp_path)], use_server_dir=False)
+    elapsed = time.time() - start_time
+
+    server_thread.join()
+
+    # Should complete after roughly 2-3 seconds (when server becomes available)
+    assert 1 < elapsed < 6.0