diff --git a/internal/agent/supervisor/supervisor.go b/internal/agent/supervisor/supervisor.go index cc605aa..60e3005 100644 --- a/internal/agent/supervisor/supervisor.go +++ b/internal/agent/supervisor/supervisor.go @@ -113,7 +113,10 @@ type Shard struct { exitCode int logs *ringBuffer - cmd *exec.Cmd + // cmdline is an adopted engine's command line, from its state file: it + // has no cmd of this agent's to read it from. + cmdline []string + cmd *exec.Cmd // sidecars are the shard's supporting processes, torn down with it. sidecars []*exec.Cmd cancel context.CancelFunc @@ -729,11 +732,17 @@ type shardState struct { // writeState records everything needed to recognise and describe the process later. func (s *Supervisor) writeState(sh *Shard) error { + // An adopted engine was started by an earlier agent: its command line + // is the one its state file kept. + cmdline := sh.cmdline + if sh.cmd != nil { + cmdline = sh.cmd.Args + } b, err := json.Marshal(shardState{ ShardID: sh.ID, ReplicaID: sh.Spec.ReplicaID, PID: sh.pid, - Cmdline: sh.cmd.Args, + Cmdline: cmdline, StartedAt: sh.started, Spec: sh.Spec, }) @@ -830,6 +839,7 @@ func (s *Supervisor) Adopt(o Orphan) (*Shard, error) { started: o.StartedAt, done: make(chan struct{}), adopted: true, + cmdline: o.Cmdline, } s.shards[o.ShardID] = sh s.mu.Unlock() diff --git a/internal/agent/supervisor/supervisor_test.go b/internal/agent/supervisor/supervisor_test.go index 5b542ef..6a89996 100644 --- a/internal/agent/supervisor/supervisor_test.go +++ b/internal/agent/supervisor/supervisor_test.go @@ -873,3 +873,35 @@ func TestWhatAShardComputesOnSurvivesTheAgent(t *testing.T) { t.Errorf("not kept: %+v", c) } } + +func TestAnAdoptedShardsComputeIsKeptWithoutCrashing(t *testing.T) { + s, _ := newSupervisor(t) + spec := engines.ShardSpec{ReplicaID: "rep-1", ShardID: "a", HTTPPort: 20100} + if _, err := s.Start(t.Context(), "a", spec, shellCommand("sleep 30")); err != nil { + t.Fatal(err) + } + // An agent restart: the engine is adopted by a new supervisor, which did + // not start it and has no command of its own for it. + s2, err := New(Config{WorkDir: s.cfg.WorkDir, StopTimeout: 2 * time.Second}, + &scriptedHealth{}, slog.New(slog.DiscardHandler)) + if err != nil { + t.Fatal(err) + } + orphans, _ := s2.FindOrphans() + if len(orphans) != 1 { + t.Fatalf("%d orphans", len(orphans)) + } + if _, err := s2.Adopt(orphans[0]); err != nil { + t.Fatal(err) + } + t.Cleanup(func() { _ = s2.Stop(context.Background(), "a") }) + s2.SetCompute("a", engines.Compute{Backend: "CPU"}) + + // And a third agent still finds it, with its command line and compute. + s3, _ := New(Config{WorkDir: s.cfg.WorkDir, StopTimeout: 2 * time.Second}, + &scriptedHealth{}, slog.New(slog.DiscardHandler)) + again, _ := s3.FindOrphans() + if len(again) != 1 || len(again[0].Cmdline) == 0 || again[0].Spec.Compute == nil { + t.Fatalf("lost: %+v", again) + } +} diff --git a/internal/cli/enginecheck.go b/internal/cli/enginecheck.go index 1bbb13b..461e07f 100644 --- a/internal/cli/enginecheck.go +++ b/internal/cli/enginecheck.go @@ -309,14 +309,30 @@ func checkDevices(ctx context.Context, name string, fb engines.CPUFallback, out return nil } -// installedLlama is the newest llama-server the agent installed here. +// llamaRoot is where the agent installs llama.cpp builds. +var llamaRoot = "/opt/llmgw/engines/llamacpp" + +// installedLlama is the llama-server the agent installed here: under +// llamaRoot//llama-/, or one level up for older +// layouts; on a machine with NVIDIA GPUs a CUDA build first, then the newest. func installedLlama() string { - m, _ := filepath.Glob("/opt/llmgw/engines/llamacpp/*/llama-server") - if len(m) == 0 { + var found []string + for _, pat := range []string{"*/llama-server", "*/*/llama-server"} { + m, _ := filepath.Glob(filepath.Join(llamaRoot, pat)) + found = append(found, m...) + } + if len(found) == 0 { return "" } - sort.Strings(m) - return m[len(m)-1] + sort.Strings(found) + if _, err := exec.LookPath("nvidia-smi"); err == nil { + for i := len(found) - 1; i >= 0; i-- { + if strings.Contains(found[i], "cuda") { + return found[i] + } + } + } + return found[len(found)-1] } // weightDirs are where a node keeps models it has fetched. diff --git a/internal/cli/enginecheck_gpu_test.go b/internal/cli/enginecheck_gpu_test.go index 61fec52..99f92c0 100644 --- a/internal/cli/enginecheck_gpu_test.go +++ b/internal/cli/enginecheck_gpu_test.go @@ -64,3 +64,23 @@ func TestEngineCheckComputesATokenOnTheGPU(t *testing.T) { t.Fatalf("a build that falls back to the CPU: %v %s", err, out.String()) } } + +func TestEngineCheckFindsTheBuildTheAgentInstalled(t *testing.T) { + old := llamaRoot + llamaRoot = t.TempDir() + t.Cleanup(func() { llamaRoot = old }) + if got := installedLlama(); got != "" { + t.Fatalf("nothing installed, found %q", got) + } + for _, p := range []string{"b10991-cpu/llama-b10991/llama-server", "b10991-cuda-13.3/llama-b10991/llama-server"} { + full := filepath.Join(llamaRoot, p) + _ = os.MkdirAll(filepath.Dir(full), 0o755) + _ = os.WriteFile(full, []byte("x"), 0o755) + } + dir := t.TempDir() + _ = os.WriteFile(filepath.Join(dir, "nvidia-smi"), []byte("#!/bin/sh\n"), 0o755) + t.Setenv("PATH", dir) + if got := installedLlama(); !strings.Contains(got, "cuda-13.3/llama-b10991/llama-server") { + t.Errorf("with GPUs, not the CUDA build: %q", got) + } +}