fix(landscape): bound Compute's per-worker memory and clean up on failure

Review round 1: size per-worker buffers off the largest actual window
instead of WindowSize^2 (avoids allocating ~200+ MiB per worker for
small/partial rasters), cap workers at 8 regardless of GOMAXPROCS, remove
partially-written outputs when a run fails, wrap input-open errors with
node/group context, guard checkGrids against a node with no inputs, and
clamp a non-positive WindowSize so raster.Windows can't loop forever.

Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
This commit is contained in:
2026-09-05 13:30:23 +05:00
parent bc524f3d98
commit ff6d5fc412
3 changed files with 209 additions and 12 deletions

View File

@@ -3,6 +3,7 @@ package landscape
import (
"context"
"fmt"
"os"
"path/filepath"
"runtime"
"sync"
@@ -15,9 +16,16 @@ import (
// Compute runs one node over its input rasters. See the package doc for the
// concurrency model: windows are processed by a worker pool; each worker owns
// its own Reader handles (GDAL datasets are not goroutine-safe) and writes are
// serialised per output. Window order does not affect the result.
// serialised per output. Window order does not affect the result. On any
// error, no output files are left behind: Compute best-effort removes every
// output it created before returning.
func Compute(ctx context.Context, rio raster.IO, node Node, inputs map[string]string, outDir string, p Params) (map[string]string, error) {
p = p.WithDefaults()
if p.WindowSize <= 0 {
// WithDefaults only fills in a zero value; guard against a negative
// one too, since raster.Windows would otherwise loop forever.
p.WindowSize = DefaultParams().WindowSize
}
paths := make([]string, len(node.Inputs))
for i, g := range node.Inputs {
path, ok := inputs[g]
@@ -35,30 +43,47 @@ func Compute(ctx context.Context, rio raster.IO, node Node, inputs map[string]st
outPaths := make(map[string]string, len(node.Outputs))
writers := make([]raster.Writer, len(node.Outputs))
createdPaths := make([]string, 0, len(node.Outputs))
for j, g := range node.Outputs {
path := filepath.Join(outDir, g+".tif")
w, err := rio.Create(ctx, path, grid, raster.DefaultWriterOptions())
if err != nil {
closeWriters(writers[:j])
removeOutputs(createdPaths)
return nil, fmt.Errorf("node %s: create output %s: %w", node.Code, g, err)
}
writers[j] = w
outPaths[g] = path
createdPaths = append(createdPaths, path)
}
windows := raster.Windows(grid.Width, grid.Height, p.WindowSize)
// The first window raster.Windows produces is always the largest: it
// starts at (0,0) so both its width and height are the un-clipped
// min(WindowSize, dimension). Size per-worker buffers off it instead of
// WindowSize*WindowSize, which would over-allocate whenever the raster is
// smaller than one window (the common case for small/test rasters, and
// for the last row/column of any raster).
var size int
if len(windows) > 0 {
size = windows[0].Width * windows[0].Height
}
jobs := make(chan raster.Window)
var writeMu sync.Mutex
eg, gctx := errgroup.WithContext(ctx)
workers := min(runtime.GOMAXPROCS(0), len(windows))
// Cap workers well below GOMAXPROCS too: each worker holds a full set of
// per-window input/output buffers, and nodes can have many inputs and
// outputs (e.g. node 20: 6 inputs, 5 outputs), so an unbounded worker
// count multiplies memory with core count and with concurrently running
// nodes.
workers := min(runtime.GOMAXPROCS(0), len(windows), 8)
for w := 0; w < workers; w++ {
eg.Go(func() error {
readers, err := openAll(gctx, rio, paths)
readers, err := openAll(gctx, rio, node, paths)
if err != nil {
return err
}
defer closeReaders(readers)
size := p.WindowSize * p.WindowSize
in := make([][]float64, len(paths))
for i := range in {
in[i] = make([]float64, size)
@@ -109,6 +134,7 @@ func Compute(ctx context.Context, rio raster.IO, node Node, inputs map[string]st
runErr = err
}
if runErr != nil {
removeOutputs(createdPaths)
return nil, runErr
}
return outPaths, nil
@@ -117,11 +143,14 @@ func Compute(ctx context.Context, rio raster.IO, node Node, inputs map[string]st
// checkGrids opens every input once, verifies all grids match the first, and
// returns that grid.
func checkGrids(ctx context.Context, rio raster.IO, node Node, paths []string) (raster.Grid, error) {
readers, err := openAll(ctx, rio, paths)
readers, err := openAll(ctx, rio, node, paths)
if err != nil {
return raster.Grid{}, err
}
defer closeReaders(readers)
if len(readers) == 0 {
return raster.Grid{}, fmt.Errorf("node %s (%s): has no inputs", node.Code, node.Name)
}
ref := readers[0].Grid()
for i := 1; i < len(readers); i++ {
if g := readers[i].Grid(); !rio.SameGrid(ref, g) {
@@ -132,13 +161,20 @@ func checkGrids(ctx context.Context, rio raster.IO, node Node, paths []string) (
return ref, nil
}
func openAll(ctx context.Context, rio raster.IO, paths []string) ([]raster.Reader, error) {
// openAll opens paths in order, wrapping any failure with the node code and
// the input group name (paths is always node.Inputs resolved to files, in
// the same order) for context.
func openAll(ctx context.Context, rio raster.IO, node Node, paths []string) ([]raster.Reader, error) {
readers := make([]raster.Reader, 0, len(paths))
for _, p := range paths {
for i, p := range paths {
r, err := rio.Open(ctx, p)
if err != nil {
closeReaders(readers)
return nil, err
name := p
if i < len(node.Inputs) {
name = node.Inputs[i]
}
return nil, fmt.Errorf("node %s: input %s: %w", node.Code, name, err)
}
readers = append(readers, r)
}
@@ -165,3 +201,11 @@ func closeWriters(ws []raster.Writer) error {
}
return first
}
// removeOutputs best-effort deletes every path in paths, ignoring errors. It
// is used to leave no partially-written output behind after a failed run.
func removeOutputs(paths []string) {
for _, p := range paths {
_ = os.Remove(p)
}
}