superfluid_daemon/

fleet_join.rs

//! `superfluid node join`: this machine serves a model to a fleet head it dials, and the head
//! takes nodes that dial in.

use std::net::{SocketAddr, TcpListener, TcpStream, ToSocketAddrs};
use std::path::{Path, PathBuf};
use std::sync::{Arc, Mutex};
use std::time::Duration;

use superfluid_linkf::{Endpoint, LinkFError, Role};

use crate::fleet::FleetHead;
use crate::fleet_manager::FleetManager;
use crate::nodeagent::NodeAgent;
use crate::runtime_pick::RuntimeId;
use crate::runtimes::Catalog;
use crate::{Daemon, EngineHost, NoCodec, SessionStore};

pub const HELP: &str = "Serve a model to a fleet head from this machine.
Usage: superfluid node join [<head-host:port>] --model <path|org/model[:tag]> [options]
  Dials the head (found on the local network when no address is given), serves the
  model to it, and dials again whenever the link drops. The head's token, from
  `superfluid fleet token` there, authenticates both ends and encrypts the link.
Options:
  --model <m>             The model; an id is pulled first, as `serve` pulls it
  --token <value>         The head's fleet token, kept at $SUPERFLUID_HOME/fleet/token
  --runtime <id>          The runtime, when the model's format does not decide it
  --max-context <N>       This node's context window [default: sized for this device]
  --max-batch <N>         This node's lanes [default: 8]
  --identity <name>       The name the head shows for this node [default: the host name]
  --model-name <id>       The name the head matches on [default: what the head announces,
                          else the model's id]
  --offline               Pull nothing; the model must be in the runtime's cache
  --park                  Keep each prompt's prefix on disk, under
                          $SUPERFLUID_HOME/fleet/nodes/<identity>/park, so a restarted node
                          starts the head's sessions warm
";

const CONNECT: Duration = Duration::from_secs(5);
const HANDSHAKE: Duration = Duration::from_secs(10);
const DISCOVER: Duration = Duration::from_secs(30);
const RETRY_MAX: Duration = Duration::from_secs(30);
const REFUSED: &str = "the head refused this node's token; copy the head's `superfluid fleet token` here with --token";
/// Peers in the handshake at once; more wait in the backlog, so a burst of connects cannot
/// take a thread each.
const HANDSHAKES_AT_ONCE: usize = 8;

pub fn run(args: &[String]) {
    match args.first().map(String::as_str) {
        Some("join") => join(&args[1..]),
        _ => {
            eprint!("{HELP}");
            std::process::exit(2);
        }
    }
}

fn fail(msg: impl std::fmt::Display) -> ! {
    eprintln!("superfluid: {msg}");
    std::process::exit(2);
}

fn join(args: &[String]) {
    // The node-agent's warnings (a fenced session, a silent head) go to stderr, as noded's do.
    let filter = std::env::var("RUST_LOG").unwrap_or_else(|_| "warn".to_string());
    if let Err(e) = crate::telemetry::init_stderr(&filter) {
        eprintln!("superfluid: logging: {e}");
    }
    let (mut head, mut model, mut token, mut runtime, mut identity, mut model_name) = (None, None, None, None, None, None);
    let (mut max_context, mut max_batch, mut offline, mut park) = (None, 8u32, false, false);
    let mut it = args.iter();
    while let Some(a) = it.next() {
        let mut val = || it.next().cloned().unwrap_or_else(|| fail(format!("{a} takes a value")));
        match a.as_str() {
            "--model" => model = Some(val()),
            "--token" => token = Some(val()),
            "--runtime" => runtime = Some(val()),
            "--identity" => identity = Some(val()),
            "--model-name" => model_name = Some(val()),
            "--max-context" => max_context = Some(val().parse().unwrap_or_else(|_| fail("--max-context takes a number"))),
            "--max-batch" => max_batch = val().parse().unwrap_or_else(|_| fail("--max-batch takes a number")),
            "--offline" => offline = true,
            "--park" => park = true,
            "-h" | "--help" => {
                print!("{HELP}");
                return;
            }
            s if s.starts_with('-') => fail(format!("unknown flag {s}")),
            s if head.is_none() => head = Some(s.to_string()),
            s => fail(format!("unexpected argument {s}")),
        }
    }
    let Some(model) = model else { fail("--model is required") };
    let home = crate::runtimes::default_home();
    if let Some(t) = &token {
        crate::fleet_token::write(&home, t.trim().as_bytes()).unwrap_or_else(|e| fail(e));
    }
    let token = crate::fleet_token::read(&home).unwrap_or_default();
    let (target, announced) = match head {
        Some(h) => (Target { given: Some(h), found: Mutex::new(Vec::new()) }, None),
        None => {
            let (addrs, served) = discover();
            (Target { given: None, found: Mutex::new(addrs) }, Some(served))
        }
    };
    let addrs = target.addrs();
    if addrs.is_empty() {
        fail(format!("cannot resolve {target}"));
    }
    if token.is_empty() && !addrs.iter().all(|a| a.ip().is_loopback()) {
        fail("no fleet token: run `superfluid fleet token` on the head and pass --token <value> here");
    }
    let identity = identity.unwrap_or_else(crate::metrics::hostname);
    let max_batch = max_batch.max(1);
    let state = NodeState::open(&home, &identity, park).unwrap_or_else(|e| fail(e));
    let loaded = load_node(&model, runtime.as_deref(), max_context, max_batch, offline, &state).unwrap_or_else(|e| fail(e));
    eprintln!("superfluid: {}", loaded.summary);
    let model_name = model_name.unwrap_or_else(|| name_among(announced.as_deref().unwrap_or_default(), loaded.model_name));
    let mut node = NodeAgent::new(Arc::new(loaded.daemon), identity.clone(), model_name.clone()).with_max_lanes(max_batch);
    if let Some(tok) = loaded.tokenizer {
        node = node.with_tokenizer_of(tok);
    }
    // A head that keeps a connection for short calls asks for it; one connection more, once.
    let (ask, asked) = std::sync::mpsc::sync_channel::<()>(1);
    let node = node.with_auth(token.clone()).with_on_control_ask(move || {
        let _ = ask.try_send(());
    });
    eprintln!("superfluid: node '{identity}' serves {model_name}; dialing the head at {target} once per lane ({max_batch})");
    // One connection per lane, each its own generation at a time on the head.
    let target = Arc::new(target);
    let dial = |lane: u32| {
        let (target, token, node) = (Arc::clone(&target), token.clone(), node.connection());
        std::thread::spawn(move || keep_dialing(&target, &token, node, lane))
    };
    let mut links: Vec<_> = (0..max_batch).map(dial).collect();
    if asked.recv().is_ok() {
        links.push(dial(max_batch));
    }
    for link in links {
        let _ = link.join();
    }
}

/// Where the head is: an address given on the command line, resolved at each dial, or the
/// addresses the local network announced, looked for again when none of them answers.
struct Target {
    given: Option<String>,
    found: Mutex<Vec<SocketAddr>>,
}

impl Target {
    fn addrs(&self) -> Vec<SocketAddr> {
        match &self.given {
            Some(h) => h.to_socket_addrs().map(|a| a.collect()).unwrap_or_default(),
            None => self.found.lock().expect("head addrs").clone(),
        }
    }

    /// After every address in `tried` failed: looks on the network again, unless the head was
    /// given or another lane already found it anew. Lanes wait on the one that is looking.
    fn refind(&self, tried: &[SocketAddr]) {
        if self.given.is_some() {
            return;
        }
        let mut found = self.found.lock().expect("head addrs");
        if *found == tried {
            if let Ok(head) = crate::fleet_mdns::find_head(DISCOVER) {
                *found = head.nodes;
            }
        }
    }
}

impl std::fmt::Display for Target {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match &self.given {
            Some(h) => f.write_str(h),
            None => match self.found.lock().expect("head addrs").first() {
                Some(a) => write!(f, "{a}"),
                None => f.write_str("the fleet head"),
            },
        }
    }
}

/// The head on the local network, however long it takes to appear.
fn discover() -> (Vec<SocketAddr>, Vec<String>) {
    eprintln!("superfluid: looking for a fleet head on the local network");
    loop {
        match crate::fleet_mdns::find_head(DISCOVER) {
            Ok(head) => {
                eprintln!("superfluid: found the head at {}, serving {}", head.nodes[0], head.models.join(", "));
                return (head.nodes, head.models);
            }
            Err(e) => eprintln!("superfluid: {e}; looking again"),
        }
    }
}

/// The name a node serves its model under: its own when the head serves a model by that name,
/// the head's one model when it serves one (the same model under another id or path), and its
/// own, said, when the head serves several and none is it.
fn name_among(served: &[String], own: String) -> String {
    match served {
        _ if served.contains(&own) => own,
        [only] => only.clone(),
        [] => own,
        _ => {
            eprintln!(
                "superfluid: the head serves {}, and none is '{own}'; give this node --model-name <one of them>",
                served.join(", ")
            );
            own
        }
    }
}

/// Dials the head and serves until the link ends, then dials again; the first lane speaks
/// for the node in the log.
fn keep_dialing(target: &Target, token: &[u8], mut conn: NodeAgent, lane: u32) {
    let say = |msg: String| {
        if lane == 0 {
            eprintln!("superfluid: {msg}");
        }
    };
    let mut backoff = Duration::from_secs(2);
    loop {
        let addrs = target.addrs();
        match dial(&addrs, token) {
            Ok(ep) => {
                say(format!("joined the head at {target}"));
                backoff = Duration::from_secs(2);
                match conn.serve(ep) {
                    Ok(()) => {
                        tracing::debug!(lane, "the link to the head closed");
                        say("the head closed the link".into())
                    }
                    Err(crate::DaemonError::Fleet(LinkFError::Unauthorized)) => fail(REFUSED),
                    Err(e) => {
                        tracing::debug!(lane, error = %e, "the link to the head ended");
                        say(format!("the link to the head ended: {e}"))
                    }
                }
            }
            Err(LinkFError::Unauthorized) => fail(REFUSED),
            Err(e) => {
                say(format!("cannot join the head at {target}: {e}"));
                target.refind(&addrs);
            }
        }
        say(format!("dialing again in {} s", backoff.as_secs()));
        std::thread::sleep(backoff);
        backoff = (backoff * 2).min(RETRY_MAX);
    }
}

pub struct LoadedNode {
    pub daemon: Daemon,
    /// The name the head sees the model under.
    pub model_name: String,
    /// The runtime, model, window and lanes, for the node's start line.
    pub summary: String,
    /// The model's tokenizer, for the head to check and, with no copy of the model, to use.
    pub tokenizer: Option<Arc<dyn superfluid_engine::Tokenizer>>,
}

/// A fleet node's model by path or id, as `superfluid serve` loads one: pulled when given by
/// id, on its runtime (installed first if it never was), in a window sized for this device
/// unless `max_context` pins it.
pub fn load_node(
    model: &str,
    runtime: Option<&str>,
    max_context: Option<i32>,
    lanes: u32,
    offline: bool,
    state: &NodeState,
) -> Result<LoadedNode, String> {
    let lanes = lanes.max(1);
    let catalog = Catalog::from_env();
    let runtime = runtime.map(|r| if r == "native" { "basert" } else { r });
    let (path, model_name, id) = resolve_model(&catalog, model, runtime, offline)?;
    let runtime = catalog.require(id).or_else(|_| crate::pulls::ensure_installed(&catalog, id, offline).map_err(String::from))?;
    let (ctx, sized) = match max_context {
        Some(n) => (n, format!("--max-context {n}")),
        None => match device_window(&runtime, &path, lanes) {
            Ok(n) => (n, "sized for this device".to_string()),
            Err(why) => (8192, format!("{why}; --max-context N pins it")),
        },
    };
    let ctx = ctx.max(512);
    let host = if runtime.worker.bin.is_file() {
        EngineHost::spawn_process(crate::WorkerSpec {
            bin: runtime.worker.bin.clone(),
            args: runtime.serve_args(&path, ctx, lanes, 0),
            stderr: None,
        })
    } else {
        let rt = crate::linked::engine(runtime.id.name())
            .ok_or_else(|| format!("the {} runtime has no worker here", runtime.id.name()))?;
        let args = crate::runtime_options::worker_args(runtime.id, &path, ctx, lanes, 0).map_err(|e| format!("model load failed: {e}"))?;
        EngineHost::try_spawn_loaded(move || rt.open(&args))
    }
    .map_err(|e| format!("model load failed: {e}"))?;
    let summary = format!(
        "runtime {} serves {}: context window {ctx} tokens ({sized}), {lanes} lanes",
        runtime.id.name(),
        path.display()
    );
    let mut node = loaded(host, Box::new(NoCodec), lanes, model_name, summary, state)?;
    node.tokenizer = crate::serve::open_tokenizer(&runtime, &path).ok();
    Ok(node)
}

/// A node on the test engine, which needs no model.
pub fn load_mock_node(max_context: Option<i32>, lanes: u32, state: &NodeState) -> Result<LoadedNode, String> {
    let lanes = lanes.max(1);
    let engine = crate::linked::engine("mock").ok_or("this build does not link the mock engine")?;
    let args = superfluid_adapter_kit::WorkerArgs {
        max_context: max_context.unwrap_or(4096).max(512),
        max_batch: lanes,
        ..Default::default()
    };
    let host = EngineHost::try_spawn_loaded(move || engine.open(&args)).map_err(|e| format!("model load failed: {e}"))?;
    loaded(host, Box::new(crate::MockCodec), lanes, "mock".into(), "the mock engine".into(), state)
}

fn loaded(
    host: EngineHost,
    codec: Box<dyn crate::TextCodec + Send + Sync>,
    lanes: u32,
    model_name: String,
    summary: String,
    state: &NodeState,
) -> Result<LoadedNode, String> {
    let store = SessionStore::open(&state.wal).map_err(|e| format!("node store: {e}"))?;
    let daemon = Daemon::with_park(store, host, codec, lanes as usize, state.park.clone());
    Ok(LoadedNode { daemon, model_name, summary, tokenizer: None })
}

/// The window this device holds for `lanes` lanes of the model, as `superfluid serve` sizes it.
fn device_window(runtime: &crate::runtimes::Runtime, model: &std::path::Path, lanes: u32) -> Result<i32, String> {
    let id = runtime.id;
    if crate::linked::sizes_context(id) {
        return crate::linked::suggest_max_context(id, &[model.to_path_buf()], &[], lanes as i32, 0)?
            .ok_or_else(|| format!("the {} runtime could not size a window for this device", id.base().name()));
    }
    let sizing = crate::runtimes::sizing_of(runtime, model)?;
    let trained = sizing.trained_context;
    let per_lane = superfluid_adapter_kit::sizing::window(&[sizing], lanes)
        .ok_or_else(|| format!("the {} runtime reported no device budget to size a window against", id.base().name()))?;
    let whole = per_lane.saturating_mul(u64::from(lanes)).min(trained.max(per_lane));
    Ok(i32::try_from(whole).unwrap_or(i32::MAX))
}

/// The node's end of the link: it dials the first address that answers, and the head still
/// opens the handshake.
fn dial(addrs: &[SocketAddr], token: &[u8]) -> Result<Endpoint, LinkFError> {
    let mut last = None;
    for addr in addrs {
        match TcpStream::connect_timeout(addr, CONNECT) {
            Ok(sock) => {
                return if token.is_empty() {
                    Endpoint::from_tcp(sock)
                } else {
                    Endpoint::from_tcp_secure(sock, token, Role::Responder, HANDSHAKE)
                };
            }
            Err(e) => last = Some(e),
        }
    }
    Err(last.unwrap_or_else(|| std::io::Error::new(std::io::ErrorKind::NotFound, "address resolved to nothing")).into())
}

pub fn loopback(addr: &str) -> bool {
    addr.starts_with("127.") || addr.starts_with("localhost:") || addr.starts_with("[::1]:")
}

/// A session store under the system temp dir, for `superfluid-noded`.
/// Where a node keeps its state, under `$SUPERFLUID_HOME/fleet/nodes/<identity>`. Its session
/// store starts empty each run: a node's sessions last one generation, and the head keeps the
/// conversations. Its parked prefixes, with `--park`, are kept across runs.
pub struct NodeState {
    pub wal: PathBuf,
    pub park: Option<PathBuf>,
}

impl NodeState {
    pub fn open(home: &Path, identity: &str, park: bool) -> Result<NodeState, String> {
        let dir = home.join("fleet").join("nodes").join(identity.replace(['/', '\\'], "_"));
        std::fs::create_dir_all(&dir).map_err(|e| format!("{}: {e}", dir.display()))?;
        let wal = dir.join("wal.log");
        for stale in [wal.clone(), dir.join("wal.log.prev")] {
            let _ = std::fs::remove_file(stale);
        }
        Ok(NodeState { wal, park: park.then(|| dir.join("park")) })
    }

    /// A session store of this process's own, and no parking.
    pub fn scratch() -> Result<NodeState, String> {
        let dir = std::env::temp_dir().join(format!("superfluid-node-{}", std::process::id()));
        std::fs::create_dir_all(&dir).map_err(|e| format!("{}: {e}", dir.display()))?;
        Ok(NodeState { wal: dir.join("wal.log"), park: None })
    }
}

/// The model's path (pulled when given by id), the name the head sees it under, and its runtime.
fn resolve_model(catalog: &Catalog, model: &str, runtime: Option<&str>, offline: bool) -> Result<(PathBuf, String, RuntimeId), String> {
    let named = |name: &str| {
        catalog.named(name).ok().or_else(|| RuntimeId::parse(name)).ok_or_else(|| format!("--runtime '{name}' is not a runtime id"))
    };
    if crate::pulls::looks_like_id(model) {
        let id = match runtime {
            Some(name) => named(name)?,
            None => RuntimeId::new(crate::cli::runtime_for_id(model)),
        };
        let flags = crate::pulls::PullFlags { offline, options: Vec::new() };
        let path = crate::pulls::pull(catalog, id, model, &flags)?;
        return Ok((path, crate::pulls::served_name(model), id));
    }
    let path = PathBuf::from(model);
    if !path.exists() {
        return Err(format!("{}: no such file or directory", path.display()));
    }
    let id = match runtime {
        Some(name) => named(name)?,
        None => catalog.reader_of(&path)?,
    };
    Ok((path.clone(), crate::registry::model_id_for(&path), id))
}

/// The head's end: every node that dials in with the token is taken into the fleet. With no
/// token, only nodes on this machine are; `model` is what the head serves, and a node that
/// serves something else is taken but warned about, since no session will go to it.
pub fn accept_nodes(listener: TcpListener, manager: Arc<FleetManager>, auth: Vec<u8>, models: Vec<String>) {
    let in_handshake = Arc::new(std::sync::atomic::AtomicUsize::new(0));
    for sock in listener.incoming() {
        let Ok(sock) = sock else { continue };
        let Ok(peer) = sock.peer_addr() else { continue };
        if auth.is_empty() && !peer.ip().is_loopback() {
            eprintln!("superfluid: a node at {peer} was refused: this head has no fleet token (`superfluid fleet token` makes one)");
            continue;
        }
        while in_handshake.load(std::sync::atomic::Ordering::Acquire) >= HANDSHAKES_AT_ONCE {
            std::thread::sleep(Duration::from_millis(20));
        }
        in_handshake.fetch_add(1, std::sync::atomic::Ordering::AcqRel);
        let (manager, auth, served, in_handshake) = (Arc::clone(&manager), auth.clone(), models.clone(), Arc::clone(&in_handshake));
        std::thread::spawn(move || {
            let opened = FleetHead::accept(sock, auth);
            in_handshake.fetch_sub(1, std::sync::atomic::Ordering::AcqRel);
            match opened {
                Ok(head) => {
                    let node = head.node();
                    let (identity, models, ctx) = (node.node_identity.clone(), node.loadable_models.join(", "), node.capacity.max_context_tokens);
                    let serves_one = node.loadable_models.iter().any(|m| served.contains(m));
                    match manager.adopt(head, peer.ip()) {
                        Ok((_, true, _)) => {
                            eprintln!("superfluid: fleet node '{identity}' joined from {peer}: {models}, context window {ctx} tokens");
                            if !serves_one {
                                eprintln!(
                                    "superfluid: fleet node '{identity}' serves {models}, none of {}; no session will go to it (set --model-name there)",
                                    served.join(", ")
                                );
                            }
                        }
                        Ok((_, false, 1)) => eprintln!("superfluid: fleet node '{identity}' rejoined from {peer}"),
                        Ok(_) => {}
                        Err(why) => eprintln!("superfluid: a node at {peer} did not join: {why}"),
                    }
                }
                Err(e) => eprintln!("superfluid: a node at {peer} did not join: {e}"),
            }
        });
    }
}

#[cfg(test)]
mod tests {
    use super::name_among;

    #[test]
    fn a_node_takes_its_own_name_when_served_else_the_heads_only_one() {
        let own = || "qwen".to_string();
        assert_eq!(name_among(&["llama".into(), "qwen".into()], own()), "qwen");
        assert_eq!(name_among(&["Qwen3-0.6B".into()], own()), "Qwen3-0.6B", "the same model under another id");
        assert_eq!(name_among(&["llama".into(), "gemma".into()], own()), "qwen", "several, none it: its own");
        assert_eq!(name_among(&[], own()), "qwen");
    }
}