diff --git a/internal/agent/supervisor/outputfile.go b/internal/agent/supervisor/outputfile.go new file mode 100644 index 0000000..ea06614 --- /dev/null +++ b/internal/agent/supervisor/outputfile.go @@ -0,0 +1,139 @@ +package supervisor + +import ( + "bytes" + "os" + "path/filepath" + "time" +) + +// An engine's output goes to a file of its own, not to a pipe of the +// agent's. A pipe has one reader, and that reader is the agent that +// started the engine: when that agent restarts, the engine it left running +// — and adopts a moment later — writes to a pipe nobody reads, and its next +// write past the pipe's buffer fails with EPIPE. A Python engine's progress +// bar then raises, and takes the request it was reporting on with it. A +// file has no reader to lose: the engine appends to it, every agent reads +// it, and the output from before a restart is there for the one after. + +// maxOutputBytes bounds one process's file: past it the file is emptied and +// the process goes on appending to it. A progress bar rewriting itself for +// an hour would otherwise fill the disk. +const maxOutputBytes = 64 << 20 + +// outputPath is where a shard's process writes. +func (s *Supervisor) outputPath(shardID, name string) string { + if name != "" { + return filepath.Join(s.cfg.WorkDir, "shard-"+shardID+"."+name+".log") + } + return filepath.Join(s.cfg.WorkDir, "shard-"+shardID+".log") +} + +// openOutput is the file a process is started with: new, for a process +// that is new. +func (s *Supervisor) openOutput(shardID, name string) (*os.File, error) { + return os.OpenFile(s.outputPath(shardID, name), os.O_CREATE|os.O_WRONLY|os.O_APPEND|os.O_TRUNC, 0o600) //nolint:gosec // the agent's own directory +} + +// follow reads a process's file into the shard's log as it is written, +// until the shard's output is stopped and the file has been read to its +// end. A file an earlier agent began is read from the start, so the output +// from before a restart is in the log too. +func (s *Supervisor) follow(sh *Shard, path, name string) { + defer sh.outDone.Done() + f, err := os.Open(path) //nolint:gosec // the agent's own directory + if err != nil { + sh.logs.add("[" + name + "] output unavailable: " + err.Error()) + return + } + defer f.Close() + var offset int64 + var rest []byte + buf := make([]byte, 64<<10) + read := func() { + for { + st, err := f.Stat() + if err != nil { + return + } + if st.Size() < offset { + // Emptied, by this agent or another: begin again. + offset, rest = 0, nil + } + if st.Size() > maxOutputBytes { + _ = os.Truncate(path, 0) + offset, rest = 0, nil + continue + } + n, err := f.ReadAt(buf, offset) + if n > 0 { + offset += int64(n) + rest = append(rest, buf[:n]...) + for { + i := bytes.IndexByte(rest, '\n') + if i < 0 { + break + } + s.addLine(sh, name, rest[:i]) + rest = rest[i+1:] + } + // A progress bar is a line rewritten with \r and never ended: + // its latest state is what is wanted, not a line a megabyte + // long. Kept, as llama.cpp and tqdm write them. + if len(rest) > 1<<20 { + s.addLine(sh, name, rest) + rest = nil + } + } + if err != nil || n == 0 { + return + } + } + } + t := time.NewTicker(200 * time.Millisecond) + defer t.Stop() + for { + select { + case <-sh.outStop: + read() + if len(rest) > 0 { + s.addLine(sh, name, rest) + } + return + case <-t.C: + read() + } + } +} + +func (s *Supervisor) addLine(sh *Shard, name string, b []byte) { + // The last state of a line rewritten with \r, as a terminal would show. + if i := bytes.LastIndexByte(b, '\r'); i >= 0 { + b = b[i+1:] + } + line := string(b) + if name != "" { + // Prefixed so the two processes of a Ray shard stay distinguishable + // in a single log tail. + line = "[" + name + "] " + line + } + sh.logs.add(line) + s.log.Debug("engine", "shard", sh.ID, "line", line) +} + +// removeOutput drops a shard's files once nothing writes to them. +func (s *Supervisor) removeOutput(shardID string) { + matches, _ := filepath.Glob(filepath.Join(s.cfg.WorkDir, "shard-"+shardID+".*log")) + for _, m := range matches { + _ = os.Remove(m) + } + _ = os.Remove(s.outputPath(shardID, "")) +} + +// stopOutput has every follower read what is left and stop. Called once +// the process has ended, before the shard is done, so what it wrote last +// is in the log by the time anyone reads the log for why it ended. +func (s *Supervisor) stopOutput(sh *Shard) { + sh.outOnce.Do(func() { close(sh.outStop) }) + sh.outDone.Wait() +} diff --git a/internal/agent/supervisor/supervisor.go b/internal/agent/supervisor/supervisor.go index 60e3005..30f1002 100644 --- a/internal/agent/supervisor/supervisor.go +++ b/internal/agent/supervisor/supervisor.go @@ -1,12 +1,10 @@ package supervisor import ( - "bufio" "context" "encoding/json" "errors" "fmt" - "io" "log/slog" "net" "os" @@ -104,6 +102,12 @@ type Shard struct { ID string Spec engines.ShardSpec + // outStop ends the followers of the processes' output files, and + // outDone is their having read to the end. + outStop chan struct{} + outOnce sync.Once + outDone sync.WaitGroup + mu sync.Mutex state State pid int @@ -248,6 +252,7 @@ func (s *Supervisor) Start(ctx context.Context, id string, spec engines.ShardSpe logs: newRingBuffer(s.cfg.LogLines), started: time.Now(), done: make(chan struct{}), + outStop: make(chan struct{}), } s.shards[id] = sh s.mu.Unlock() @@ -315,14 +320,15 @@ func (s *Supervisor) launch(ctx context.Context, sh *Shard, cmds []engines.Comma // CommandContext kills only the leader, so replace that with a group kill. c.Cancel = func() error { return killGroup(c.Process.Pid, syscall.SIGKILL) } - stdout, err := c.StdoutPipe() + out, err := s.openOutput(sh.ID, "") if err != nil { cancel() return err } - c.Stderr = c.Stdout.(io.Writer) // interleave, as the engine writes them - - if err := c.Start(); err != nil { + c.Stdout, c.Stderr = out, out // interleaved, as the engine writes them + err = c.Start() + out.Close() // the process holds its own + if err != nil { cancel() return fmt.Errorf("starting %s: %w", cmd.Path, err) } @@ -339,7 +345,8 @@ func (s *Supervisor) launch(ctx context.Context, sh *Shard, cmds []engines.Comma "shard", sh.ID, "err", err) } - go s.pump(sh, stdout, cmd.Name) + sh.outDone.Add(1) + go s.follow(sh, s.outputPath(sh.ID, ""), cmd.Name) go s.wait(sh) return nil } @@ -353,16 +360,18 @@ func (s *Supervisor) spawn(ctx context.Context, sh *Shard, c engines.Command) (* s.dropTo(proc) proc.Cancel = func() error { return killGroup(proc.Process.Pid, syscall.SIGKILL) } - stdout, err := proc.StdoutPipe() + out, err := s.openOutput(sh.ID, c.Name) if err != nil { return nil, err } - proc.Stderr = proc.Stdout.(io.Writer) - - if err := proc.Start(); err != nil { + proc.Stdout, proc.Stderr = out, out + err = proc.Start() + out.Close() + if err != nil { return nil, fmt.Errorf("starting %s: %w", c.Path, err) } - go s.pump(sh, stdout, c.Name) + sh.outDone.Add(1) + go s.follow(sh, s.outputPath(sh.ID, c.Name), c.Name) // Reaped so it does not become a zombie; its exit does not end the shard. go func() { _ = proc.Wait() }() return proc, nil @@ -395,24 +404,6 @@ func waitForPort(ctx context.Context, addr string, timeout time.Duration) (err e } // pump copies a process's output into the ring buffer and the log. -func (s *Supervisor) pump(sh *Shard, r io.Reader, name string) { - sc := bufio.NewScanner(r) - // Engine tracebacks and progress bars produce long lines; the default - // 64 KiB limit truncates them exactly when they matter. - sc.Buffer(make([]byte, 0, 64<<10), 1<<20) - - for sc.Scan() { - line := sc.Text() - if name != "" { - // Prefixed so the two processes of a Ray shard stay distinguishable - // in a single log tail. - line = "[" + name + "] " + line - } - sh.logs.add(line) - s.log.Debug("engine", "shard", sh.ID, "line", line) - } -} - // wait reaps the process and records why it ended. func (s *Supervisor) wait(sh *Shard) { defer close(sh.done) @@ -440,6 +431,7 @@ func (s *Supervisor) wait(sh *Shard) { final, code := sh.state, sh.exitCode sh.mu.Unlock() + s.stopOutput(sh) s.removeState(sh) if final == Failed { s.log.Error("shard failed", "shard", sh.ID, "exit_code", code, "err", err, @@ -760,6 +752,7 @@ func (s *Supervisor) writeState(sh *Shard) error { func (s *Supervisor) removeState(sh *Shard) { _ = os.Remove(s.statePath(sh.ID)) + s.removeOutput(sh.ID) } // Orphan is an engine process left behind by a previous agent. @@ -838,15 +831,18 @@ func (s *Supervisor) Adopt(o Orphan) (*Shard, error) { logs: newRingBuffer(s.cfg.LogLines), started: o.StartedAt, done: make(chan struct{}), + outStop: make(chan struct{}), adopted: true, cmdline: o.Cmdline, } s.shards[o.ShardID] = sh s.mu.Unlock() - // Output is gone with the previous agent; say so rather than presenting an - // empty log as if the engine had been silent. - sh.logs.add("[adopted from a previous agent; earlier output is not available]") + // Its output file is read from the start: what it wrote under the + // previous agent is there, and what it writes now goes on arriving. + sh.logs.add("[adopted from a previous agent]") + sh.outDone.Add(1) + go s.follow(sh, s.outputPath(sh.ID, ""), "") go s.poll(sh) return sh, nil @@ -874,6 +870,7 @@ func (s *Supervisor) poll(sh *Shard) { } sh.mu.Unlock() + s.stopOutput(sh) s.removeState(sh) s.log.Info("adopted shard exited", "shard", sh.ID, "expected", expected) return @@ -900,6 +897,7 @@ func (s *Supervisor) KillOrphan(o Orphan) error { } // An already-missing state file is success: the previous agent may have // reaped the process itself between the scan and this call. + s.removeOutput(o.ShardID) if err := os.Remove(s.statePath(o.ShardID)); err != nil && !errors.Is(err, os.ErrNotExist) { return err } diff --git a/internal/agent/supervisor/supervisor_test.go b/internal/agent/supervisor/supervisor_test.go index 6a89996..54cd9c1 100644 --- a/internal/agent/supervisor/supervisor_test.go +++ b/internal/agent/supervisor/supervisor_test.go @@ -1,6 +1,7 @@ package supervisor import ( + "bytes" "context" "errors" "fmt" @@ -905,3 +906,78 @@ func TestAnAdoptedShardsComputeIsKeptWithoutCrashing(t *testing.T) { t.Fatalf("lost: %+v", again) } } + +func TestAnAdoptedEngineGoesOnWritingAndItsEarlierOutputIsKept(t *testing.T) { + // An engine's output went to a pipe of the agent's. The agent that + // restarted and adopted the engine left that pipe with no reader, and + // the engine's next write past the pipe's buffer failed with EPIPE — + // which, in a Python engine, raised inside whatever was being reported + // on and failed the request. Output goes to a file now: a reader is not + // something a file can lose, and the lines from before the restart are + // there for the agent after it. + s, _ := newSupervisor(t) + // 256 KiB at once outgrows a pipe's buffer; a reader gone means EPIPE. + sh, err := s.Start(t.Context(), "a", engines.ShardSpec{ReplicaID: "r"}, + shellCommand(`echo before-restart; sleep 1; head -c 262144 /dev/zero | tr '\0' x; echo; echo after-restart-done; sleep 30`)) + if err != nil { + t.Fatal(err) + } + eventually(t, 3*time.Second, "the first line was never captured", func() bool { + return strings.Contains(strings.Join(sh.Logs(10), "\n"), "before-restart") + }) + // The agent goes away without stopping the engine: nothing reads the + // output between here and the adoption. + pid := sh.Status().PID + sh.outOnce.Do(func() { close(sh.outStop) }) + sh.outDone.Wait() + + s2, err := New(Config{WorkDir: s.cfg.WorkDir, StopTimeout: 3 * time.Second, LogLines: 64}, + &scriptedHealth{}, slog.New(slog.DiscardHandler)) + if err != nil { + t.Fatal(err) + } + orphans, _ := s2.FindOrphans() + adopted, err := s2.Adopt(orphans[0]) + if err != nil { + t.Fatal(err) + } + eventually(t, 5*time.Second, "the engine's output after the restart never arrived: "+strings.Join(adopted.Logs(5), " | "), func() bool { + return strings.Contains(strings.Join(adopted.Logs(64), "\n"), "after-restart-done") + }) + if !strings.Contains(strings.Join(adopted.Logs(64), "\n"), "before-restart") { + t.Error("the output from before the restart is not in the adopted log") + } + if !processAlive(pid) { + t.Fatal("the engine died writing output nobody was reading") + } + if err := s2.Stop(t.Context(), "a"); err != nil { + t.Fatal(err) + } + if _, err := os.Stat(s2.outputPath("a", "")); !errors.Is(err, os.ErrNotExist) { + t.Errorf("the output file was not removed with the shard: %v", err) + } +} + +func TestAnOutputFileIsEmptiedPastItsBound(t *testing.T) { + s, _ := newSupervisor(t) + path := s.outputPath("a", "") + // A file already past the bound, as a long-running engine leaves one. + big, _ := os.OpenFile(path, os.O_CREATE|os.O_WRONLY, 0o644) + _, _ = big.Write(bytes.Repeat([]byte("x"), maxOutputBytes+1)) + big.Close() + sh := &Shard{ID: "a", logs: newRingBuffer(8), done: make(chan struct{}), outStop: make(chan struct{})} + sh.outDone.Add(1) + go s.follow(sh, path, "") + time.Sleep(300 * time.Millisecond) + f, _ := os.OpenFile(path, os.O_WRONLY|os.O_APPEND, 0o644) + _, _ = f.WriteString("still here\n") + f.Close() + s.stopOutput(sh) + st, _ := os.Stat(path) + if st.Size() > 1024 { + t.Errorf("the file was not emptied: %d bytes", st.Size()) + } + if !strings.Contains(strings.Join(sh.Logs(8), "\n"), "still here") { + t.Errorf("a line written after the file was emptied was lost: %v", sh.Logs(8)) + } +}