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:
@@ -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)
|
||||
}
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user