Stage 7a-i: TP worker lifecycle scaffolding
All checks were successful
CI / Format (push) Successful in 36s
build-prerelease / Resolve version stamps (push) Successful in 39s
CI / Clippy (push) Successful in 2m12s
CI / Test (push) Successful in 4m25s
CI / Build cortex SRPM (push) Has been skipped
CI / Build neuron SRPM (push) Has been skipped
CI / Publish cortex to COPR (push) Has been skipped
CI / Publish neuron to COPR (push) Has been skipped
CI / Bump version in source (push) Has been skipped
build-prerelease / Build neuron-blackwell (push) Successful in 3m49s
build-prerelease / Build cortex binary (push) Successful in 4m22s
build-prerelease / Package cortex RPM (push) Successful in 1m23s
build-prerelease / Build neuron-ampere (push) Successful in 5m9s
build-prerelease / Build neuron-ada (push) Successful in 4m59s
build-prerelease / Package helexa-neuron-ada RPM (push) Successful in 2m53s
build-prerelease / Package helexa-neuron-ampere RPM (push) Successful in 2m59s
build-prerelease / Package helexa-neuron-blackwell RPM (push) Successful in 3m38s
build-prerelease / Publish to rpm.lair.cafe (unstable) (push) Successful in 1m8s
All checks were successful
CI / Format (push) Successful in 36s
build-prerelease / Resolve version stamps (push) Successful in 39s
CI / Clippy (push) Successful in 2m12s
CI / Test (push) Successful in 4m25s
CI / Build cortex SRPM (push) Has been skipped
CI / Build neuron SRPM (push) Has been skipped
CI / Publish cortex to COPR (push) Has been skipped
CI / Publish neuron to COPR (push) Has been skipped
CI / Bump version in source (push) Has been skipped
build-prerelease / Build neuron-blackwell (push) Successful in 3m49s
build-prerelease / Build cortex binary (push) Successful in 4m22s
build-prerelease / Package cortex RPM (push) Successful in 1m23s
build-prerelease / Build neuron-ampere (push) Successful in 5m9s
build-prerelease / Build neuron-ada (push) Successful in 4m59s
build-prerelease / Package helexa-neuron-ada RPM (push) Successful in 2m53s
build-prerelease / Package helexa-neuron-ampere RPM (push) Successful in 2m59s
build-prerelease / Package helexa-neuron-blackwell RPM (push) Successful in 3m38s
build-prerelease / Publish to rpm.lair.cafe (unstable) (push) Successful in 1m8s
Leader → worker process plumbing for tensor parallelism. The neuron
binary picks up two modes: default (the existing daemon, axum + HTTP)
and `--worker` (a bare RPC loop driven over stdin/stdout). The leader
spawns one worker per non-zero NCCL rank via tokio::process::Command
on the same binary path (production: /proc/self/exe; tests:
env!("CARGO_BIN_EXE_neuron")) and talks to each over newline-
delimited JSON.
Protocol (harness/tp/rpc.rs) is serde-tagged from the start —
WorkerRequest::{Ping, Init, NcclSanityCheck, Shutdown} and
WorkerResponse::{Pong, InitOk, NcclSanityResult, Bye, Error}, both
`#[serde(tag = "op", rename_all = "snake_case")]`. Adding ops in 7b/7c
is purely additive; unknown ops on the wire fail to parse (verified
in unit tests).
7a-i scope:
- WorkerPool::spawn(binary, world_size, devices) forks ranks 1..N as
subprocesses, captures stdin/stdout, kills on drop.
- ping_all() round-trips a Ping to every worker and validates the
returned rank.
- shutdown() sends Shutdown to each worker, awaits Bye, reaps.
- Worker mode: parse Ping/Shutdown, return Pong/Bye; Init and
NcclSanityCheck return Error{kind="not_implemented_7a_i"} so a 7a-ii
binary speaking the same wire is a drop-in replacement (the kind
field signals "real NCCL lands in the next commit").
- CandleHarness::load_model refuses tensor_parallel > 1 with a clear
message until 7b is in.
Three integration tests in tests/tp_worker_lifecycle.rs cover spawn/
ping/shutdown for 2- and 3-worker pools, plus the
not_implemented_7a_i contract test for Init. Seven rpc serde unit
tests assert the wire shape (op tags, field names, unknown-op
rejection). All pass on the dev host; no CUDA required.
Stage 7a-ii (next): the real NCCL Comm::from_rank wiring behind the
existing Init/NcclSanityCheck op surface, CUDA-gated. Verifiable on
beast's 2×5090.
Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
This commit is contained in:
@@ -1,21 +1,52 @@
|
||||
use anyhow::Result;
|
||||
use clap::Parser;
|
||||
use neuron::{api, config::NeuronConfig, discovery, harness::HarnessRegistry, health, startup};
|
||||
use neuron::{
|
||||
api,
|
||||
config::NeuronConfig,
|
||||
discovery,
|
||||
harness::{HarnessRegistry, tp},
|
||||
health, startup,
|
||||
};
|
||||
use std::sync::Arc;
|
||||
use std::time::Instant;
|
||||
use tokio::sync::RwLock;
|
||||
use tracing_subscriber::EnvFilter;
|
||||
|
||||
/// Top-level CLI. The same binary runs as either the public neuron
|
||||
/// daemon (default) or a tensor-parallel worker subprocess (when
|
||||
/// `--worker` is set) spawned by the leader on the same host.
|
||||
#[derive(Parser)]
|
||||
#[command(name = "neuron")]
|
||||
#[command(about = "Per-node daemon for cortex inference clusters")]
|
||||
#[command(version)]
|
||||
struct Args {
|
||||
/// Port to listen on (overrides config file).
|
||||
/// Run in tensor-parallel worker mode. The leader process spawns
|
||||
/// one of these per non-zero NCCL rank and drives it over
|
||||
/// newline-delimited JSON on stdin/stdout. Worker mode skips
|
||||
/// discovery, the HTTP listener, and the health poller — it's a
|
||||
/// pure RPC loop.
|
||||
#[arg(long, default_value_t = false)]
|
||||
worker: bool,
|
||||
|
||||
/// NCCL rank for worker mode. Ignored when `--worker` is not set.
|
||||
#[arg(long, default_value_t = 0)]
|
||||
rank: u32,
|
||||
|
||||
/// Total NCCL world size for worker mode. Ignored when `--worker`
|
||||
/// is not set.
|
||||
#[arg(long, default_value_t = 1)]
|
||||
tp_size: u32,
|
||||
|
||||
/// CUDA device index for worker mode. Ignored when `--worker` is
|
||||
/// not set.
|
||||
#[arg(long, default_value_t = 0)]
|
||||
cuda_device: u32,
|
||||
|
||||
/// Port to listen on (overrides config file). Daemon mode only.
|
||||
#[arg(short, long)]
|
||||
port: Option<u16>,
|
||||
|
||||
/// Path to the neuron config file.
|
||||
/// Path to the neuron config file. Daemon mode only.
|
||||
#[arg(short, long, default_value = "neuron.toml")]
|
||||
config: String,
|
||||
}
|
||||
@@ -23,6 +54,7 @@ struct Args {
|
||||
#[tokio::main]
|
||||
async fn main() -> Result<()> {
|
||||
tracing_subscriber::fmt()
|
||||
.with_writer(std::io::stderr)
|
||||
.with_env_filter(
|
||||
EnvFilter::try_from_default_env()
|
||||
.unwrap_or_else(|_| EnvFilter::new("info,neuron=debug")),
|
||||
@@ -31,6 +63,19 @@ async fn main() -> Result<()> {
|
||||
|
||||
let args = Args::parse();
|
||||
|
||||
if args.worker {
|
||||
return tp::worker::run(tp::worker::WorkerConfig {
|
||||
rank: args.rank,
|
||||
world_size: args.tp_size,
|
||||
cuda_device: args.cuda_device,
|
||||
})
|
||||
.await;
|
||||
}
|
||||
|
||||
daemon(args).await
|
||||
}
|
||||
|
||||
async fn daemon(args: Args) -> Result<()> {
|
||||
let cfg = NeuronConfig::load(&args.config).unwrap_or_else(|e| {
|
||||
tracing::warn!(path = %args.config, error = %e, "config not found, using defaults");
|
||||
NeuronConfig::default()
|
||||
|
||||
Reference in New Issue
Block a user