llmgw/internal/agent/supervisor/outputfile.go

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()
}