// Package cluster runs LocalAI as real child processes for end-to-end tests. // // The in-process suites cannot express frontend-replica failure: there is no // process to kill, no second replica to race, and no real HTTP boundary between // a worker and the frontend it registered with. This package starts the same // binary an operator runs, one process per frontend replica and one per worker, // against containerised Postgres. Postgres is the only container: a distributed // deployment needs it and the frontends' own HTTP listener, and nothing else. package cluster import ( "fmt" "math/rand/v2" "net" "net/http" "os" "os/exec" "path/filepath" "strconv" "strings" "syscall" "time" "github.com/mudler/LocalAI/pkg/httpclient" "github.com/phayes/freeport" ) // StaleBusURL is the address of a broker that is not running, and is not meant // to be. // // LOCALAI_NATS_URL is still accepted and ignored by every process this harness // starts, so an operator's existing command line, unit file or Helm values file // starts unchanged after the broker is shut down. This suite covers that // promise by handing over THIS value: a value pointing at a live server would // let a regression that dialled it pass unnoticed, while a dead one turns the // same regression into a startup failure in every cluster spec. // // It is a constant here rather than an Options field, because there is nothing // left for a caller to choose: no process reads the variable, so a per-cluster // value would only be a knob that changes nothing. const StaleBusURL = "nats://127.0.0.1:1" // Options configures a cluster. Every field without a default is required. type Options struct { // Binary is the path to a built local-ai. Binary string // MockBackend is the path to the mock-backend binary. When set it is copied // into each worker's backends directory as "mock-backend", which is the name // model YAML refers to (see tests/e2e/e2e_suite_test.go:75). MockBackend string // MockBackendAliases installs the same fixture backend under additional // production backend names. This lets endpoint conformance retain the real // backend-specific capability discovery and validation path. MockBackendAliases []string // PGDSN points at infrastructure the caller already started. // // There is no NatsURL beside it any more. Removing the FIELD rather than // ignoring it means a spec that still sets one fails to compile, which is // the only way a harness option stops being set by accident. PGDSN string // LogDir receives one file per process. Never empty: a cluster failure is // unreadable without them. LogDir string RegistrationToken string // default "e2e-token" AdminEmail string // default "admin@e2e.local" // RequireNodeApproval keeps new workers pending until an administrator // approves them. The default remains auto-approval for the existing cluster // matrix; authentication conformance opts into the production approval flow. RequireNodeApproval bool // DistributedRequireAuth turns on the same fail-closed umbrella setting on // every frontend and worker process. It is separate from RegistrationToken: // the ordinary harness has always supplied a token, while this option proves // the binaries were explicitly told that an empty token is fatal. DistributedRequireAuth bool Frontends int Workers int // AgentWorkers is how many `local-ai agent-worker` processes to start // alongside the backend workers. // // They exist so a spec can hold the two kinds of worker side by side in one // cluster. An agent worker holds a tunnel of its own and is reached through // it and through nothing else, so it gets a real node_connections row whose // departure really does age past the grace and it is subject to every // tunnel-departure rule a backend worker is. A spec that asserted only on // backend workers could not tell "the rule reaches both fleets" from // "nothing here looks at agent workers". // // They register through the same WorkerFrontendURL hook as backend workers, // so a spec that puts a balancer in front of the fleet gets one for its // agent workers too. Without that, killing a replica would orphan the agent // worker's heartbeats and it would go unhealthy for a reason that has // nothing to do with what the spec is about. AgentWorkers int // ReconnectGrace sets LOCALAI_WORKER_RECONNECT_GRACE on every frontend: how // long a worker whose tunnel was lost is read as reconnecting rather than // gone. Zero leaves the binary's own default (90s). // // It is an Option rather than a per-spec environment edit because two // specs need it pulled in OPPOSITE directions and neither can use the // default: one has to prove nothing is reaped inside the window and needs // it longer than a re-home takes, and one has to prove a departed worker // stops being reported healthy and would otherwise wait 90 seconds to say // so. ReconnectGrace time.Duration // SpreadWorkerRegistrations sends worker i to frontend i%Frontends instead // of sending every worker to frontend 0. // // Off by default, and deliberately so: the cross-replica session specs read // a node at frontend 1 that only frontend 0 was ever told about, and // spreading registrations would leave them passing while proving nothing. // It exists for the racing-replicas spec, which is about two replicas // writing to one roster concurrently and cannot express that at all while // every worker registers through the same process. SpreadWorkerRegistrations bool // Models is written into every frontend's models directory before that // replica starts, keyed by file name. It is how a spec gets a model // configuration in front of the frontend at all: the models directory is // scanned at startup, so a file written afterwards is not guaranteed to be // seen, and there is no admin endpoint that creates a config. // // Frontends only. A worker is handed model artifacts by the frontend's file // staging, over the tunnel, and pre-seeding the worker would hide whether // that worked. Models map[string]string // WorkerFrontendURL rewrites the URL worker i registers and holds its // tunnel against. It is called once per worker, after every frontend is // serving, and is given the URL the worker would otherwise have been handed // plus every frontend's URL in index order, so a hook can put a proxy or a // load balancer in front of one replica or of all of them. // // It exists for two things the fixed per-replica URL cannot express. One is // a worker that survives its replica: LOCALAI_REGISTER_TO is resolved once // at boot and is the tunnel endpoint as well as the registration one, so a // worker pointed straight at a replica has nowhere to reconnect to when // that replica dies, and the re-home this feature is built on cannot // happen. The other is the suite's negative control, which needs a worker // that registers, heartbeats and reports healthy exactly as usual while its // tunnel dial never reaches a frontend; nothing else can produce that, // because LOCALAI_WORKER_TUNNEL=false is refused at startup and a worker // that never started proves nothing about a worker reachable some other // way. WorkerFrontendURL func(worker int, registrar string, frontends []string) string // Galleries is the model-gallery list every frontend runs with, in the JSON // form --galleries takes. Empty leaves the binary's own default list, which // is what every spec that does not name this option keeps. // // It exists for the fan-out specs, which need the thing a gallery operation // blocks on to be something the spec can hold open. A model install's first // act is to fetch its gallery's index, and the gallery worker runs one // operation at a time, so an index served by a gate the spec controls is // what turns "the queued broadcast went out before the terminal one" from a // timing coincidence into an arrangement. // // Setting it also turns the startup estimate warmer OFF, and that is not // tidiness. The warmer fetches the same index in the background at boot, so // with it on the process-wide index cache can be filled by something other // than the operation the spec is holding, and the operation then never // blocks at all. The spec would still pass, on an ordering nothing enforced. Galleries string // ConformanceStaging pins the worker's ephemeral staging capacity while // keeping admission enforcement enabled. It is intentionally opt-in: only // fixture-heavy conformance deployments should override production's // filesystem-derived limits. ConformanceStaging bool } // Process is one running local-ai. type Process struct { Name string // Cmd is exposed for signalling only. Never call Cmd.Wait on it: the reaper // goroutine started in spawn owns it, a second Wait races the first, and // waitErr is only safe to read after <-p.exited. Cmd *exec.Cmd Port int LogPath string logFile *os.File // exited closes once the reaper has collected the process. Only the reaper // calls Cmd.Wait, so nothing else may: a second Wait on the same Cmd races // the first and corrupts ProcessState. // // A closed exited proves the child is gone. An open one proves nothing: it // is still open for the whole interval between the child exiting and waitid // collecting it, during which the child is a zombie that signal 0 reports as // alive. Anything asserting on a process being dead must poll, not sample. exited chan struct{} waitErr error } // Cluster is a running set of frontend and worker processes. type Cluster struct { opts Options frontends []*Process workers []*Process // agentWorkers are kept apart from workers rather than appended to it. Every // index-taking method on this type means "backend worker i", and folding the // two together would silently renumber them for every existing spec. agentWorkers []*Process baseDir string } const ( defaultRegistrationToken = "e2e-token" defaultAdminEmail = "admin@e2e.local" // The failover specs have bounded observation windows sized around a 60s // stale threshold. Production deliberately defaults to 5m, so the harness // pins its shorter test clock explicitly. Checkpoint below the worker's 10s // heartbeat interval so every beat remains durable and the 60s threshold // measures worker silence rather than checkpoint lag. testStaleNodeThreshold = "60s" testNodeHeartbeatCheckpoint = "5s" // testHMACSecret is shared by every frontend so a session minted at one // replica validates at all of them. See the note in startFrontend. // #nosec G101 -- a fixed secret for a throwaway cluster the suite starts // and kills; it never leaves the test process. testHMACSecret = "e2e-cluster-hmac-secret" readinessTimeout = 90 * time.Second readinessPoll = 200 * time.Millisecond // processGracefulExitTimeout lets the real binaries run the same shutdown // handlers they run under an orchestrator. In particular, workers must stop // their backend process groups before exiting. It exceeds processmanager's // 15s backend grace, but remains bounded so one stuck child cannot consume a // suite-wide Ginkgo timeout. processGracefulExitTimeout = 20 * time.Second processKillExitTimeout = 10 * time.Second ) func (o *Options) applyDefaults() { if o.RegistrationToken == "" { o.RegistrationToken = defaultRegistrationToken } if o.AdminEmail == "" { o.AdminEmail = defaultAdminEmail } } func (o Options) validate() error { if o.Frontends < 1 { return fmt.Errorf("cluster needs at least one frontend, got %d", o.Frontends) } if o.LogDir == "" { return fmt.Errorf("cluster needs a LogDir: process logs are the only way to read a cluster failure") } if st, err := os.Stat(o.Binary); err != nil || st.IsDir() { return fmt.Errorf("local-ai binary not found at %q (run: make build)", o.Binary) } return nil } func workerStagingEnv(opts Options, workerRoot string) []string { if !opts.ConformanceStaging { return nil } return []string{ "LOCALAI_EPHEMERAL_STAGING_BYTE_LIMIT=1073741824", "LOCALAI_EPHEMERAL_STAGING_MIN_FREE_BYTES=1", "LOCALAI_MOCK_EXPECT_STAGING_ROOT=" + workerRoot, } } func authenticatedDeploymentEnv(opts Options) []string { env := []string{fmt.Sprintf("LOCALAI_AUTO_APPROVE_NODES=%t", !opts.RequireNodeApproval)} if opts.DistributedRequireAuth { env = append(env, "LOCALAI_DISTRIBUTED_REQUIRE_AUTH=true") } return env } func environmentWithout(env []string, key string) []string { prefix := key + "=" filtered := make([]string, 0, len(env)) for _, entry := range env { if !strings.HasPrefix(entry, prefix) { filtered = append(filtered, entry) } } return filtered } // Start brings up the cluster. It blocks until every frontend answers /readyz // and every worker process has been spawned. It does NOT wait for workers to // register: that needs an authenticated admin session, so a caller that depends // on registration must poll /api/nodes itself. func Start(opts Options) (*Cluster, error) { opts.applyDefaults() if err := opts.validate(); err != nil { return nil, err } baseDir, err := os.MkdirTemp("", "localai-cluster-*") if err != nil { return nil, fmt.Errorf("creating cluster work dir: %w", err) } c := &Cluster{opts: opts, baseDir: baseDir} for i := 0; i < opts.Frontends; i++ { p, err := c.startFrontend(i, 0) if err != nil { c.Stop() return nil, err } c.frontends = append(c.frontends, p) } for i := 0; i < opts.Workers; i++ { p, err := c.startWorker(i) if err != nil { c.Stop() return nil, err } c.workers = append(c.workers, p) } for i := 0; i < opts.AgentWorkers; i++ { p, err := c.startAgentWorker(i) if err != nil { c.Stop() return nil, err } c.agentWorkers = append(c.agentWorkers, p) } return c, nil } // startFrontend starts frontend i. A port <= 0 allocates a fresh one; a pinned // port exists for restart: workers take LOCALAI_REGISTER_TO once at boot and // never re-resolve it, so a replica that comes back on a new port is // unreachable by exactly the workers that registered with it. func (c *Cluster) startFrontend(i int, port int) (*Process, error) { if port <= 0 { allocated, err := freeport.GetFreePort() if err != nil { return nil, fmt.Errorf("allocating frontend port: %w", err) } port = allocated } name := frontendName(i) dir := c.frontendDir(i) if err := os.MkdirAll(filepath.Join(dir, "models"), 0o750); err != nil { return nil, fmt.Errorf("creating %s dirs: %w", name, err) } if err := os.MkdirAll(filepath.Join(dir, "backends"), 0o750); err != nil { return nil, fmt.Errorf("creating %s dirs: %w", name, err) } // Without an explicit LOCALAI_DATA_PATH every child resolves DataPath to // ${cwd}/data (core/cli/run.go:48), which under `go test` is inside the // source tree and shared by every replica: one collectiondb, one task and // job store for processes that are meant to be independent. dataPath := c.frontendDataDir(i) if err := os.MkdirAll(dataPath, 0o750); err != nil { return nil, fmt.Errorf("creating %s dirs: %w", name, err) } for file, content := range c.opts.Models { path := filepath.Join(dir, "models", file) if err := os.WriteFile(path, []byte(content), 0o600); err != nil { return nil, fmt.Errorf("writing %s for %s: %w", file, name, err) } } // #nosec G204 -- the binary is the one this suite built and the args are // the suite's own, not user input. cmd := exec.Command(c.opts.Binary, "run", "--address", fmt.Sprintf("127.0.0.1:%d", port), "--models-path", filepath.Join(dir, "models"), "--backends-path", filepath.Join(dir, "backends"), ) // Cmd.Environ() is the parent environment this Cmd would already run with; // the children need PATH, HOME and the Go/CI environment intact. cmd.Env = append(environmentWithout(cmd.Environ(), "LOCALAI_API_TOKEN"), "LOCALAI_DISTRIBUTED=true", // Deliberately handed a broker URL, and deliberately a dead one. This // is the shape of an operator's existing unit file on the day they shut // the broker down, and the promise it covers is that such a command // line still STARTS. A live address would let a regression that dialled // it pass; this one turns such a regression into a startup failure in // every cluster spec. "LOCALAI_NATS_URL="+StaleBusURL, "LOCALAI_AUTH=true", "LOCALAI_AUTH_DATABASE_URL="+c.opts.PGDSN, "LOCALAI_ADMIN_EMAIL="+c.opts.AdminEmail, "LOCALAI_DATA_PATH="+dataPath, // Session rows are keyed by HMAC-SHA256(token, APIKeyHMACSecret), and // the secret is generated per instance into {DataPath}/.hmac_secret // unless pinned (core/application/startup.go:141-148). Now that each // replica owns its data directory, an unpinned secret would differ per // replica, so the cookie minted at frontend 0 would hash to a session // row that does not exist at frontend 1 and every post-failover // /api/nodes call would 401 with nothing in the logs to explain it. // Pinning makes the cross-replica session a property of the harness. "LOCALAI_AUTH_HMAC_SECRET="+testHMACSecret, "LOCALAI_REGISTRATION_TOKEN="+c.opts.RegistrationToken, "LOCALAI_STALE_NODE_THRESHOLD="+testStaleNodeThreshold, "LOCALAI_NODE_HEARTBEAT_CHECKPOINT="+testNodeHeartbeatCheckpoint, // Every replica here shares one host, so the address a peer dials is // this process's own loopback address. It has to be said explicitly: // the automatic discovery asks which local address routes to // PostgreSQL, and this suite's PostgreSQL is a container published on // 127.0.0.1, so the discovery refuses (correctly) rather than // advertising a loopback address that would mean "yourself" on a // multi-host deployment. fmt.Sprintf("LOCALAI_DISTRIBUTED_ADVERTISE_ADDR=127.0.0.1:%d", port), "DEBUG=true", ) cmd.Env = append(cmd.Env, authenticatedDeploymentEnv(c.opts)...) if c.opts.ReconnectGrace > 0 { cmd.Env = append(cmd.Env, "LOCALAI_WORKER_RECONNECT_GRACE="+c.opts.ReconnectGrace.String()) } if c.opts.Galleries != "" { // Both, together, or neither: see Options.Galleries for why the warmer // has to be off wherever the gallery index is a spec's gate. cmd.Env = append(cmd.Env, "LOCALAI_GALLERIES="+c.opts.Galleries, "LOCALAI_AUTOLOAD_GALLERIES=false") } p, err := c.spawn(name, cmd, port) if err != nil { return nil, err } if err := waitReady(p, fmt.Sprintf("http://127.0.0.1:%d/readyz", port)); err != nil { // The caller never sees this process, so Stop() will never reach it: // reap it here or it outlives the suite holding a port and a log handle. p.terminate() return nil, fmt.Errorf("%s never became ready (see %s): %w", name, p.LogPath, err) } return p, nil } func (c *Cluster) startWorker(i int) (*Process, error) { // One contiguous block, laid out the way production lays it out. See // reserveWorkerPorts. grpcPort, err := reserveWorkerPorts() if err != nil { return nil, fmt.Errorf("allocating worker ports: %w", err) } httpPort := grpcPort - 1 maxPort := grpcPort + workerPortBlockSize - 1 name := fmt.Sprintf("worker-%d", i) dir := filepath.Join(c.baseDir, name) backends := filepath.Join(dir, "backends") if err := os.MkdirAll(filepath.Join(dir, "models"), 0o750); err != nil { return nil, fmt.Errorf("creating %s dirs: %w", name, err) } if err := os.MkdirAll(backends, 0o750); err != nil { return nil, fmt.Errorf("creating %s dirs: %w", name, err) } if c.opts.MockBackend != "" { if err := copyExecutable(c.opts.MockBackend, filepath.Join(backends, "mock-backend")); err != nil { return nil, fmt.Errorf("installing mock backend for %s: %w", name, err) } for _, alias := range c.opts.MockBackendAliases { if filepath.Base(alias) != alias || alias == "." || alias == "" { return nil, fmt.Errorf("invalid mock backend alias %q", alias) } if err := copyExecutable(c.opts.MockBackend, filepath.Join(backends, alias)); err != nil { return nil, fmt.Errorf("installing mock backend alias %q for %s: %w", alias, name, err) } } } // #nosec G204 -- the binary is the one this suite built and the args are // the suite's own, not user input. cmd := exec.Command(c.opts.Binary, "worker", "--models-path", filepath.Join(dir, "models"), "--backends-path", backends, ) cmd.Env = append(environmentWithout(cmd.Environ(), "LOCALAI_API_TOKEN"), // Ports only. A worker advertises nothing, so there is no advertise // address to set; these exist to keep concurrently running workers off // each other's ports, not to make anything reachable. Every bind is // loopback whatever is set here. // // The max port is what keeps the backend allocator inside the block // reserved for this worker. Without it the allocator walks upward to // 65535 (core/services/worker/registration.go, effectiveMaxPort), so a // worker running enough backends walks straight out of its block and // into whatever else this host is using. fmt.Sprintf("LOCALAI_SERVE_ADDR=127.0.0.1:%d", grpcPort), fmt.Sprintf("LOCALAI_HTTP_ADDR=127.0.0.1:%d", httpPort), fmt.Sprintf("LOCALAI_GRPC_MAX_PORT=%d", maxPort), // Workers register with frontend 0 ONLY unless the caller opts into // SpreadWorkerRegistrations, and the cross-replica session specs depend // on that default. They prove a session minted at frontend 0 resolves at // frontend 1 by reading a node that only frontend 0 was ever told about; // register the worker everywhere and they still pass while proving // nothing. // // Nothing in those specs can detect the change. The registry keys nodes // by name and preserves ids across the shared Postgres // (core/services/nodes/registry.go:522-527), so a roster read at // frontend 1 looks identical either way. Anyone changing the default // here must revisit tests/e2e/distributed/cluster_baseline_test.go by // hand. // // The registrar also fixes where this worker's heartbeats go for the // rest of its life: the loop posts to the URL it was given at boot and // never re-resolves it (core/cli/workerregistry/client.go), so killing a // worker's registrar orphans that worker rather than failing it over. // // No LOCALAI_NATS_URL: a backend worker connects to no bus, and passing // one would make every spec here prove the tunnel-only path works while // quietly handing the worker the thing it is supposed to do without. // The frontends above still get it. "LOCALAI_REGISTER_TO="+c.workerFrontendURL(i), "LOCALAI_NODE_NAME="+name, "LOCALAI_REGISTRATION_TOKEN="+c.opts.RegistrationToken, "DEBUG=true", ) if c.opts.DistributedRequireAuth { cmd.Env = append(cmd.Env, "LOCALAI_DISTRIBUTED_REQUIRE_AUTH=true") } cmd.Env = append(cmd.Env, workerStagingEnv(c.opts, dir)...) return c.spawn(name, cmd, grpcPort) } // startAgentWorker starts agent worker i. // // It is `local-ai agent-worker`, not `local-ai worker`, and the difference is // the whole point of having it here: an agent worker runs no backend processes, // so it is the second worker KIND every rule about a departed tunnel now has to // hold for. It dials a tunnel of its own, but binds only loopback for the // control plane behind it, so there is still no port to reserve and no // reachable readiness endpoint to wait on; a spec learns it is up by finding it // in the roster. func (c *Cluster) startAgentWorker(i int) (*Process, error) { name := agentWorkerName(i) // #nosec G204 -- the binary is the one this suite built and the args are // the suite's own, not user input. cmd := exec.Command(c.opts.Binary, "agent-worker") cmd.Env = append(environmentWithout(cmd.Environ(), "LOCALAI_API_TOKEN"), // Still set, and still ignored, on the same terms as the frontend's // above: this suite keeps covering the promise that an operator's // existing --nats-url does not break an agent worker. Nothing in the // process reads it. "LOCALAI_NATS_URL="+StaleBusURL, "LOCALAI_REGISTER_TO="+c.workerFrontendURL(i), "LOCALAI_NODE_NAME="+name, "LOCALAI_REGISTRATION_TOKEN="+c.opts.RegistrationToken, "DEBUG=true", ) if c.opts.DistributedRequireAuth { cmd.Env = append(cmd.Env, "LOCALAI_DISTRIBUTED_REQUIRE_AUTH=true") } return c.spawn(name, cmd, 0) } // ProcKind names which of a cluster's three process families an environment // read is about. type ProcKind string const ( ProcFrontend ProcKind = "frontend" ProcWorker ProcKind = "worker" ProcAgentWorker ProcKind = "agent-worker" ) // ProcessEnviron reads the LIVE environment of process i of kind k from // /proc//environ. // // Generalised from WorkerEnviron, which an earlier phase added for exactly this // assertion and which could only see workers. Reading /proc rather than Cmd.Env // is the whole point: Cmd.Env is what the harness INTENDED, and an assertion // about intent cannot fail when a variable arrives from the parent environment, // a profile or a wrapper. // // Linux only, like the rest of this package. A platform without /proc returns // the read error rather than falling back, so the assertion fails loudly // instead of quietly becoming the weaker one. func (c *Cluster) ProcessEnviron(k ProcKind, i int) ([]string, error) { switch k { case ProcFrontend: if err := c.checkFrontendIndex(i); err != nil { return nil, err } return processEnviron(c.frontends[i]) case ProcWorker: if err := c.checkWorkerIndex(i); err != nil { return nil, err } return processEnviron(c.workers[i]) case ProcAgentWorker: if i < 0 || i >= len(c.agentWorkers) { return nil, fmt.Errorf("agent worker %d out of range (cluster has %d)", i, len(c.agentWorkers)) } return processEnviron(c.agentWorkers[i]) default: return nil, fmt.Errorf("unknown process kind %q", k) } } // WorkerEnviron is the environment of worker i's RUNNING PROCESS, read from // /proc. // // Not Cmd.Env, deliberately. A spec asserting that a worker runs with no bus // URL is asserting about the process, and Cmd.Env is the harness telling the // spec what the harness meant to do: the two agree by construction, so a spec // reading it proves the harness consistent with itself and nothing about the // binary. /proc//environ is what the kernel handed the process. // // Kept as a named wrapper rather than replaced by ProcessEnviron at its call // sites: cluster_control_test.go calls it and re-aiming a passing spec for a // rename is churn that hides the new coverage. func (c *Cluster) WorkerEnviron(i int) ([]string, error) { return c.ProcessEnviron(ProcWorker, i) } // FrontendEnviron is the environment of frontend i's RUNNING PROCESS, read the // same way. // // It exists so a spec proving a WORKER runs with no broker URL can show that // the deployment it joined was handed one: without that half, a harness which // simply stopped setting the variable anywhere would satisfy the worker // assertion just as well. Read from /proc rather than from the harness options // for the same reason the worker's is: it is a fact about the process. func (c *Cluster) FrontendEnviron(i int) ([]string, error) { return c.ProcessEnviron(ProcFrontend, i) } // processEnviron reads a running child's environment out of /proc. func processEnviron(p *Process) ([]string, error) { if p == nil || p.Cmd == nil || p.Cmd.Process == nil { return nil, fmt.Errorf("process is not running") } raw, err := os.ReadFile(fmt.Sprintf("/proc/%d/environ", p.Cmd.Process.Pid)) if err != nil { return nil, fmt.Errorf("reading the environment of %s from /proc: %w", p.Name, err) } // NUL separated, with a trailing NUL on a non-empty environment. entries := strings.Split(string(raw), "\x00") out := make([]string, 0, len(entries)) for _, e := range entries { if e != "" { out = append(out, e) } } return out, nil } // AgentWorkerName is the node name agent worker i registered under. func (c *Cluster) AgentWorkerName(i int) string { if i < 0 || i >= len(c.agentWorkers) { return "" } return c.agentWorkers[i].Name } func agentWorkerName(i int) string { return fmt.Sprintf("agent-worker-%d", i) } const ( // workerPortBlockSize is how many ports one worker reserves: one for its // HTTP file-transfer server and the rest for backend processes. A spec that // loads more models than this on one worker exhausts the allocator, which // fails the backend start by name (ErrNoFreePort) instead of colliding. workerPortBlockSize = 24 // Workers take their ports from BELOW the ephemeral range, which on Linux // starts at 32768 by default. That is not tidiness. Ports the kernel hands // out for outbound connections are exactly the ports a long-lived process // full of outbound connections is liable to be holding when a backend tries // to bind one, and this suite's workers hold a tunnel and a registration // client each. workerPortFloor = 20000 workerPortCeiling = 31000 // workerPortAttempts bounds the search for a free block before giving up. workerPortAttempts = 200 ) // reserveWorkerPorts returns the base gRPC port of a contiguous run of ports // nothing is currently listening on, laid out the way a real worker lays them // out: the HTTP file-transfer server at base-1, and backend processes upward // from base. // // A contiguous block rather than two independent freeport allocations, and the // difference is a defect this suite actually hit. The allocator hands backend // processes basePort, basePort+1, basePort+2 and so on, and it used to do that // from its own bookkeeping alone, so the moment freeport returned two ADJACENT // ports the second backend started on a worker was handed the worker's own HTTP // server's port and died at startup with "address already in use". freeport // returns adjacent ports often, and no spec started two backends on one worker // until the tunnel load measurement did, so it presented as a spec that failed // about one run in three. // // core/services/worker/supervisor.go, allocatePort, now probes every candidate // before handing it out, which removes that collision at the other end too. // This block is kept for the reasons that survive the fix: it keeps the whole // range out of the ports the kernel allocates from, so the probe has nothing to // skip, and it bounds the allocator to a known block so an exhausted range in // this suite is a bug in this suite. // // Neither is race free and neither can be: both probes close each listener // before anything binds it. Together they are the difference between // "sometimes" and "not observed". func reserveWorkerPorts() (int, error) { for attempt := 0; attempt < workerPortAttempts; attempt++ { // #nosec G404 -- picks a candidate port block to probe. The probe, not // the draw, decides whether the block is usable. base := workerPortFloor + rand.IntN(workerPortCeiling-workerPortFloor) if blockIsFree(base-1, workerPortBlockSize+1) { return base, nil } } return 0, fmt.Errorf("no free run of %d ports in [%d, %d) after %d attempts", workerPortBlockSize+1, workerPortFloor, workerPortCeiling, workerPortAttempts) } // blockIsFree reports whether count ports from first can all be bound on // loopback right now. Every listener is held until the whole run is proven, so // a run is not accepted on the strength of one port that a previous iteration // of this same loop had just released. func blockIsFree(first, count int) bool { held := make([]net.Listener, 0, count) defer func() { for _, l := range held { _ = l.Close() } }() for port := first; port < first+count; port++ { l, err := net.Listen("tcp", fmt.Sprintf("127.0.0.1:%d", port)) if err != nil { return false } held = append(held, l) } return true } // workerFrontendURL is the URL worker i is told to register and tunnel // against: its registrar's, unless the caller installed a hook. func (c *Cluster) workerFrontendURL(i int) string { registrar := c.FrontendURL(c.registrarFor(i)) if c.opts.WorkerFrontendURL == nil { return registrar } frontends := make([]string, 0, len(c.frontends)) for index := range c.frontends { frontends = append(frontends, c.FrontendURL(index)) } return c.opts.WorkerFrontendURL(i, registrar, frontends) } func (c *Cluster) spawn(name string, cmd *exec.Cmd, port int) (*Process, error) { logPath := filepath.Join(c.opts.LogDir, name+".log") // Append rather than truncate: a restarted process reopens the same path, and // the log of the instance that died is the one a failover post-mortem needs. // #nosec G304 -- logPath is built from the suite's own temp dir and the // process name it just started, not user input. f, err := os.OpenFile(logPath, os.O_CREATE|os.O_WRONLY|os.O_APPEND, 0o600) if err != nil { return nil, fmt.Errorf("creating log file for %s: %w", name, err) } cmd.Stdout = f cmd.Stderr = f if err := cmd.Start(); err != nil { _ = f.Close() return nil, fmt.Errorf("starting %s: %w", name, err) } p := &Process{Name: name, Cmd: cmd, Port: port, LogPath: logPath, logFile: f, exited: make(chan struct{})} // One reaper per process, joined by terminate(): a child that dies on its own // is collected immediately, so waiters learn about it instead of polling a // dead port until the readiness timeout. go func() { p.waitErr = cmd.Wait() close(p.exited) }() return p, nil } // terminate asks the process to shut down, waits for the reaper, and falls back // to SIGKILL after a bounded grace. The graceful first step is load-bearing: // local-ai workers use their SIGTERM handler to stop backend process groups, // while killing only the worker would orphan those children. Safe to call more // than once and on a process that already exited. func (p *Process) terminate() { if p == nil || p.Cmd == nil || p.Cmd.Process == nil { return } select { case <-p.exited: if p.logFile != nil { _ = p.logFile.Close() } return default: } _ = p.Cmd.Process.Signal(syscall.SIGTERM) select { case <-p.exited: if p.logFile != nil { _ = p.logFile.Close() } return case <-time.After(processGracefulExitTimeout): } _ = p.Cmd.Process.Kill() select { case <-p.exited: case <-time.After(processKillExitTimeout): fmt.Printf("warning: %s did not exit within %s after SIGKILL; continuing teardown\n", p.Name, processKillExitTimeout) } if p.logFile != nil { _ = p.logFile.Close() } } // FrontendURL is the base URL of frontend i. func (c *Cluster) FrontendURL(i int) string { return fmt.Sprintf("http://127.0.0.1:%d", c.frontends[i].Port) } // RegistrationToken is the shared secret this cluster was started with. It // authenticates worker registration AND the replica-to-replica peer link, so a // spec acting as a peer needs it rather than a second literal that can drift // from Options. func (c *Cluster) RegistrationToken() string { return c.opts.RegistrationToken } // FrontendBackendsDir is the directory frontend i installs its OWN backends // into. // // It is exported for one assertion, and a filesystem one rather than an API // one: a spec proving a node backend listing came from the WORKER has to show // the frontend that answered does not have that backend itself, and the // /backends endpoint cannot say so, because in distributed mode it reports the // cluster's backends rather than this process's. func (c *Cluster) FrontendBackendsDir(i int) (string, error) { if err := c.checkFrontendIndex(i); err != nil { return "", err } return filepath.Join(c.frontendDir(i), "backends"), nil } // WorkerModelsDir returns worker i's private models directory. It is exposed // for binary conformance assertions that must distinguish a frontend-only // seed from the copy produced by distributed file staging. func (c *Cluster) WorkerModelsDir(i int) (string, error) { if err := c.checkWorkerIndex(i); err != nil { return "", err } return filepath.Join(c.baseDir, c.workers[i].Name, "models"), nil } // WorkerBackendPIDs returns the live backend children currently owned by // worker i whose executable name matches backend. Reading the kernel's process // tree, instead of inferring a load from API state, lets teardown assertions // retain the exact processes which must disappear after the worker exits. // // This helper is Linux-only, like ProcessEnviron below: the binary cluster // suite runs on Linux and deliberately fails rather than weakening a lifecycle // assertion on platforms without /proc. func (c *Cluster) WorkerBackendPIDs(i int, backend string) ([]int, error) { if err := c.checkWorkerIndex(i); err != nil { return nil, err } workerPID := c.workers[i].Cmd.Process.Pid entries, err := os.ReadDir("/proc") if err != nil { return nil, fmt.Errorf("reading process table for worker %d: %w", i, err) } var matches []int for _, entry := range entries { pid, err := strconv.Atoi(entry.Name()) if err != nil { continue } status, err := os.ReadFile(fmt.Sprintf("/proc/%d/status", pid)) if err != nil { continue // the process exited between readdir and read } if !strings.Contains(string(status), fmt.Sprintf("\nPPid:\t%d\n", workerPID)) { continue } executable, err := os.Readlink(fmt.Sprintf("/proc/%d/exe", pid)) if err == nil && filepath.Base(executable) == backend { matches = append(matches, pid) } } return matches, nil } // WorkerName is the node name worker i registered under. func (c *Cluster) WorkerName(i int) string { return c.workers[i].Name } // registrarFor is the frontend index worker i registers and heartbeats with. // // It is read at spawn time and baked into the worker's environment, so it is // also the answer to "which replica's death orphans this worker". func (c *Cluster) registrarFor(worker int) int { if !c.opts.SpreadWorkerRegistrations || c.opts.Frontends < 1 { return 0 } return worker % c.opts.Frontends } // WorkerRegistrar is registrarFor, exported so a spec can say which replica it // is about to kill relative to a worker instead of re-deriving the rule. // // It returns an error rather than indexing blindly, like every other exported // method here that takes an index. Gomega treats the trailing error as one that // must be nil, so Expect(c.WorkerRegistrar(0)).To(...) reads unchanged at the // call site while an out-of-range index fails the spec by name instead of // silently answering 0, which is a real frontend index and would send a spec // off to kill the wrong replica. func (c *Cluster) WorkerRegistrar(worker int) (int, error) { if err := c.checkWorkerIndex(worker); err != nil { return 0, err } return c.registrarFor(worker), nil } // Stop terminates every process and removes the work directory. Logs survive in // LogDir, which the caller owns. func (c *Cluster) Stop() { // Start returns (nil, err) after stopping itself, so a spec that defers // c.Stop before asserting the error would otherwise nil-deref. if c == nil { return } for _, p := range append(append(append([]*Process{}, c.workers...), c.agentWorkers...), c.frontends...) { p.terminate() } if c.baseDir != "" { _ = os.RemoveAll(c.baseDir) } } // DumpLogs writes every process log to stdout. Call from an AfterEach guarded by // CurrentSpecReport().Failed(). func (c *Cluster) DumpLogs() { for _, p := range append(append(append([]*Process{}, c.frontends...), c.workers...), c.agentWorkers...) { if p == nil { continue } data, err := os.ReadFile(p.LogPath) if err != nil { fmt.Printf("=== %s: log unreadable: %v\n", p.Name, err) continue } fmt.Printf("=== %s (%s) ===\n%s\n", p.Name, p.LogPath, string(data)) } } func waitReady(p *Process, url string) error { deadline := time.Now().Add(readinessTimeout) client := httpclient.NewWithTimeout(2 * time.Second) var last error for time.Now().Before(deadline) { select { case <-p.exited: if p.waitErr == nil { return fmt.Errorf("process exited cleanly before becoming ready") } return fmt.Errorf("process exited before becoming ready: %w", p.waitErr) default: } resp, err := client.Get(url) if err == nil { _ = resp.Body.Close() if resp.StatusCode == http.StatusOK { return nil } last = fmt.Errorf("status %d", resp.StatusCode) } else { last = err } time.Sleep(readinessPoll) } return fmt.Errorf("not ready within %s: %w", readinessTimeout, last) } func copyExecutable(src, dst string) error { // #nosec G304 -- src is a path the suite built itself, not user input. data, err := os.ReadFile(src) if err != nil { return fmt.Errorf("reading %s: %w", src, err) } // #nosec G306,G703 -- dst is a binary the suite is about to execute, so // it has to carry the execute bit, and both paths are the suite's own. if err := os.WriteFile(dst, data, 0o755); err != nil { return fmt.Errorf("writing %s: %w", dst, err) } return nil }