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 (
|
import (
|
||||||
"context"
|
"context"
|
||||||
"fmt"
|
"fmt"
|
||||||
|
"os"
|
||||||
"path/filepath"
|
"path/filepath"
|
||||||
"runtime"
|
"runtime"
|
||||||
"sync"
|
"sync"
|
||||||
@@ -15,9 +16,16 @@ import (
|
|||||||
// Compute runs one node over its input rasters. See the package doc for the
|
// 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
|
// 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
|
// 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) {
|
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()
|
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))
|
paths := make([]string, len(node.Inputs))
|
||||||
for i, g := range node.Inputs {
|
for i, g := range node.Inputs {
|
||||||
path, ok := inputs[g]
|
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))
|
outPaths := make(map[string]string, len(node.Outputs))
|
||||||
writers := make([]raster.Writer, len(node.Outputs))
|
writers := make([]raster.Writer, len(node.Outputs))
|
||||||
|
createdPaths := make([]string, 0, len(node.Outputs))
|
||||||
for j, g := range node.Outputs {
|
for j, g := range node.Outputs {
|
||||||
path := filepath.Join(outDir, g+".tif")
|
path := filepath.Join(outDir, g+".tif")
|
||||||
w, err := rio.Create(ctx, path, grid, raster.DefaultWriterOptions())
|
w, err := rio.Create(ctx, path, grid, raster.DefaultWriterOptions())
|
||||||
if err != nil {
|
if err != nil {
|
||||||
closeWriters(writers[:j])
|
closeWriters(writers[:j])
|
||||||
|
removeOutputs(createdPaths)
|
||||||
return nil, fmt.Errorf("node %s: create output %s: %w", node.Code, g, err)
|
return nil, fmt.Errorf("node %s: create output %s: %w", node.Code, g, err)
|
||||||
}
|
}
|
||||||
writers[j] = w
|
writers[j] = w
|
||||||
outPaths[g] = path
|
outPaths[g] = path
|
||||||
|
createdPaths = append(createdPaths, path)
|
||||||
}
|
}
|
||||||
|
|
||||||
windows := raster.Windows(grid.Width, grid.Height, p.WindowSize)
|
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)
|
jobs := make(chan raster.Window)
|
||||||
var writeMu sync.Mutex
|
var writeMu sync.Mutex
|
||||||
eg, gctx := errgroup.WithContext(ctx)
|
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++ {
|
for w := 0; w < workers; w++ {
|
||||||
eg.Go(func() error {
|
eg.Go(func() error {
|
||||||
readers, err := openAll(gctx, rio, paths)
|
readers, err := openAll(gctx, rio, node, paths)
|
||||||
if err != nil {
|
if err != nil {
|
||||||
return err
|
return err
|
||||||
}
|
}
|
||||||
defer closeReaders(readers)
|
defer closeReaders(readers)
|
||||||
size := p.WindowSize * p.WindowSize
|
|
||||||
in := make([][]float64, len(paths))
|
in := make([][]float64, len(paths))
|
||||||
for i := range in {
|
for i := range in {
|
||||||
in[i] = make([]float64, size)
|
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
|
runErr = err
|
||||||
}
|
}
|
||||||
if runErr != nil {
|
if runErr != nil {
|
||||||
|
removeOutputs(createdPaths)
|
||||||
return nil, runErr
|
return nil, runErr
|
||||||
}
|
}
|
||||||
return outPaths, nil
|
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
|
// checkGrids opens every input once, verifies all grids match the first, and
|
||||||
// returns that grid.
|
// returns that grid.
|
||||||
func checkGrids(ctx context.Context, rio raster.IO, node Node, paths []string) (raster.Grid, error) {
|
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 {
|
if err != nil {
|
||||||
return raster.Grid{}, err
|
return raster.Grid{}, err
|
||||||
}
|
}
|
||||||
defer closeReaders(readers)
|
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()
|
ref := readers[0].Grid()
|
||||||
for i := 1; i < len(readers); i++ {
|
for i := 1; i < len(readers); i++ {
|
||||||
if g := readers[i].Grid(); !rio.SameGrid(ref, g) {
|
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
|
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))
|
readers := make([]raster.Reader, 0, len(paths))
|
||||||
for _, p := range paths {
|
for i, p := range paths {
|
||||||
r, err := rio.Open(ctx, p)
|
r, err := rio.Open(ctx, p)
|
||||||
if err != nil {
|
if err != nil {
|
||||||
closeReaders(readers)
|
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)
|
readers = append(readers, r)
|
||||||
}
|
}
|
||||||
@@ -165,3 +201,11 @@ func closeWriters(ws []raster.Writer) error {
|
|||||||
}
|
}
|
||||||
return first
|
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)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|||||||
@@ -2,6 +2,8 @@ package landscape
|
|||||||
|
|
||||||
import (
|
import (
|
||||||
"context"
|
"context"
|
||||||
|
"fmt"
|
||||||
|
"os"
|
||||||
"path/filepath"
|
"path/filepath"
|
||||||
"testing"
|
"testing"
|
||||||
|
|
||||||
@@ -43,6 +45,9 @@ func TestCompute_E_TilesWindowsAndWritesEveryPixel(t *testing.T) {
|
|||||||
if res.grid != g {
|
if res.grid != g {
|
||||||
t.Fatal("output grid must copy the first input's grid")
|
t.Fatal("output grid must copy the first input's grid")
|
||||||
}
|
}
|
||||||
|
if res.opts != raster.DefaultWriterOptions() {
|
||||||
|
t.Fatalf("output created with %+v, want raster.DefaultWriterOptions() = %+v", res.opts, raster.DefaultWriterOptions())
|
||||||
|
}
|
||||||
for i, v := range res.data {
|
for i, v := range res.data {
|
||||||
want := 4.0
|
want := 4.0
|
||||||
if i == 0 {
|
if i == 0 {
|
||||||
@@ -99,3 +104,116 @@ func TestCompute_MultiOutputNode(t *testing.T) {
|
|||||||
t.Fatal("min(P, E) not applied per output")
|
t.Fatal("min(P, E) not applied per output")
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
// TestCompute_MultiOutputNode_TiledWindows runs the 6-input/5-output node 20
|
||||||
|
// with a window size much smaller than the raster, so per-worker output
|
||||||
|
// buffers are reused across windows of differing pixel counts (partial
|
||||||
|
// windows at the raster's right/bottom edges).
|
||||||
|
func TestCompute_MultiOutputNode_TiledWindows(t *testing.T) {
|
||||||
|
fio := newFakeIO()
|
||||||
|
w, h := 5, 3
|
||||||
|
g := grid(w, h)
|
||||||
|
paths := map[string]string{}
|
||||||
|
for i, s := range []string{SlotPIrrigation, SlotPSnow, SlotPCrop, SlotPPasture, SlotPTourism} {
|
||||||
|
paths[s] = s + ".tif"
|
||||||
|
fio.put(paths[s], g, constRaster(w, h, float64(i+1)))
|
||||||
|
}
|
||||||
|
paths[OutE] = "e.tif"
|
||||||
|
fio.put("e.tif", g, constRaster(w, h, 3))
|
||||||
|
node, _ := NodeByCode("20")
|
||||||
|
p := DefaultParams()
|
||||||
|
p.WindowSize = 2 // 5x3 with 2x2 tiles: several differently-sized windows
|
||||||
|
out, err := Compute(context.Background(), fio, node, paths, "/o2", p)
|
||||||
|
if err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
codes := []string{OutCIrrigation, OutCSnow, OutCCrop, OutCPasture, OutCTourism}
|
||||||
|
for i, code := range codes {
|
||||||
|
want := min(i+1, 3) // min(P_sector, E=3)
|
||||||
|
res := fio.files[out[code]]
|
||||||
|
for px, v := range res.data {
|
||||||
|
if v != float64(want) {
|
||||||
|
t.Fatalf("%s pixel %d = %v want %v", code, px, v, want)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// TestCompute_ZeroParamsUsesDefaults exercises Compute with a completely
|
||||||
|
// zero-valued Params, relying on WithDefaults to fill in WindowSize and the
|
||||||
|
// other node-independent knobs.
|
||||||
|
func TestCompute_ZeroParamsUsesDefaults(t *testing.T) {
|
||||||
|
fio := newFakeIO()
|
||||||
|
g := grid(4, 4)
|
||||||
|
fio.put("z.tif", g, constRaster(4, 4, 2))
|
||||||
|
fio.put("u.tif", g, constRaster(4, 4, 3))
|
||||||
|
node, _ := NodeByCode("10")
|
||||||
|
out, err := Compute(context.Background(), fio, node, map[string]string{OutZIntegrated: "z.tif", OutUIntegrated: "u.tif"}, "/z", Params{})
|
||||||
|
if err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
for _, v := range fio.files[out[OutE]].data {
|
||||||
|
if v != 3 {
|
||||||
|
t.Fatalf("pixel = %v want 3", v)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// TestCompute_NegativeWindowSizeClampedToDefault guards against
|
||||||
|
// raster.Windows looping forever on a non-positive window size: WithDefaults
|
||||||
|
// only fills in a zero value, so Compute must clamp negatives itself.
|
||||||
|
func TestCompute_NegativeWindowSizeClampedToDefault(t *testing.T) {
|
||||||
|
fio := newFakeIO()
|
||||||
|
g := grid(2, 2)
|
||||||
|
fio.put("z.tif", g, constRaster(2, 2, 2))
|
||||||
|
fio.put("u.tif", g, constRaster(2, 2, 3))
|
||||||
|
node, _ := NodeByCode("10")
|
||||||
|
p := DefaultParams()
|
||||||
|
p.WindowSize = -5
|
||||||
|
out, err := Compute(context.Background(), fio, node, map[string]string{OutZIntegrated: "z.tif", OutUIntegrated: "u.tif"}, "/neg", p)
|
||||||
|
if err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
if fio.files[out[OutE]].data[0] != 3 {
|
||||||
|
t.Fatal("unexpected result")
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// TestCompute_FailureMidRunReturnsError simulates a read failure partway
|
||||||
|
// through a run (the fake fails every ReadFloat64 against u.tif from the
|
||||||
|
// first call on) and checks Compute surfaces the error instead of returning
|
||||||
|
// a partial result.
|
||||||
|
func TestCompute_FailureMidRunReturnsError(t *testing.T) {
|
||||||
|
fio := newFakeIO()
|
||||||
|
g := grid(10, 7)
|
||||||
|
fio.put("z.tif", g, constRaster(10, 7, 2))
|
||||||
|
fio.put("u.tif", g, constRaster(10, 7, 4))
|
||||||
|
fio.failReadAfterN("u.tif", 0) // fail starting from the very first read
|
||||||
|
node, _ := NodeByCode("10")
|
||||||
|
p := DefaultParams()
|
||||||
|
p.WindowSize = 4
|
||||||
|
if _, err := Compute(context.Background(), fio, node, map[string]string{OutZIntegrated: "z.tif", OutUIntegrated: "u.tif"}, "/out-fail", p); err == nil {
|
||||||
|
t.Fatal("expected error from simulated read failure")
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// TestRemoveOutputs_DeletesFiles unit-tests the cleanup helper Compute uses
|
||||||
|
// after a failed run, against the real filesystem (the fake IO above has no
|
||||||
|
// on-disk files to assert against).
|
||||||
|
func TestRemoveOutputs_DeletesFiles(t *testing.T) {
|
||||||
|
dir := t.TempDir()
|
||||||
|
var paths []string
|
||||||
|
for i := 0; i < 3; i++ {
|
||||||
|
p := filepath.Join(dir, fmt.Sprintf("out%d.tif", i))
|
||||||
|
if err := os.WriteFile(p, []byte("data"), 0o644); err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
paths = append(paths, p)
|
||||||
|
}
|
||||||
|
removeOutputs(paths)
|
||||||
|
for _, p := range paths {
|
||||||
|
if _, err := os.Stat(p); !os.IsNotExist(err) {
|
||||||
|
t.Fatalf("expected %s to be removed, stat err = %v", p, err)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|||||||
@@ -12,11 +12,19 @@ import (
|
|||||||
type fakeIO struct {
|
type fakeIO struct {
|
||||||
mu sync.Mutex
|
mu sync.Mutex
|
||||||
files map[string]*fakeRaster
|
files map[string]*fakeRaster
|
||||||
|
|
||||||
|
// failReadAfter, when set for a path, makes every ReadFloat64 call
|
||||||
|
// against that path fail once the cumulative read count (across all
|
||||||
|
// readers/workers) exceeds the configured threshold. Used to simulate a
|
||||||
|
// mid-run I/O failure.
|
||||||
|
failReadAfter map[string]int
|
||||||
|
readCounts map[string]int
|
||||||
}
|
}
|
||||||
|
|
||||||
type fakeRaster struct {
|
type fakeRaster struct {
|
||||||
grid raster.Grid
|
grid raster.Grid
|
||||||
data []float64 // row-major
|
data []float64 // row-major
|
||||||
|
opts raster.WriterOptions
|
||||||
}
|
}
|
||||||
|
|
||||||
func newFakeIO() *fakeIO { return &fakeIO{files: map[string]*fakeRaster{}} }
|
func newFakeIO() *fakeIO { return &fakeIO{files: map[string]*fakeRaster{}} }
|
||||||
@@ -25,12 +33,39 @@ func (f *fakeIO) put(path string, g raster.Grid, data []float64) {
|
|||||||
f.files[path] = &fakeRaster{grid: g, data: data}
|
f.files[path] = &fakeRaster{grid: g, data: data}
|
||||||
}
|
}
|
||||||
|
|
||||||
type fakeReader struct{ r *fakeRaster }
|
// failReadAfterN makes ReadFloat64 against path fail starting from the
|
||||||
|
// (n+1)-th call to it.
|
||||||
|
func (f *fakeIO) failReadAfterN(path string, n int) {
|
||||||
|
f.mu.Lock()
|
||||||
|
defer f.mu.Unlock()
|
||||||
|
if f.failReadAfter == nil {
|
||||||
|
f.failReadAfter = map[string]int{}
|
||||||
|
f.readCounts = map[string]int{}
|
||||||
|
}
|
||||||
|
f.failReadAfter[path] = n
|
||||||
|
}
|
||||||
|
|
||||||
|
type fakeReader struct {
|
||||||
|
f *fakeIO
|
||||||
|
path string
|
||||||
|
r *fakeRaster
|
||||||
|
}
|
||||||
|
|
||||||
func (r *fakeReader) Grid() raster.Grid { return r.r.grid }
|
func (r *fakeReader) Grid() raster.Grid { return r.r.grid }
|
||||||
func (r *fakeReader) NoData() (float64, bool) { return -9999, true }
|
func (r *fakeReader) NoData() (float64, bool) { return -9999, true }
|
||||||
func (r *fakeReader) Close() error { return nil }
|
func (r *fakeReader) Close() error { return nil }
|
||||||
func (r *fakeReader) ReadFloat64(w raster.Window, dst []float64) error {
|
func (r *fakeReader) ReadFloat64(w raster.Window, dst []float64) error {
|
||||||
|
if r.f != nil {
|
||||||
|
r.f.mu.Lock()
|
||||||
|
if n, ok := r.f.failReadAfter[r.path]; ok {
|
||||||
|
r.f.readCounts[r.path]++
|
||||||
|
if r.f.readCounts[r.path] > n {
|
||||||
|
r.f.mu.Unlock()
|
||||||
|
return fmt.Errorf("fake: simulated read failure for %s", r.path)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
r.f.mu.Unlock()
|
||||||
|
}
|
||||||
for y := 0; y < w.Height; y++ {
|
for y := 0; y < w.Height; y++ {
|
||||||
copy(dst[y*w.Width:(y+1)*w.Width], r.r.data[(w.Y+y)*r.r.grid.Width+w.X:(w.Y+y)*r.r.grid.Width+w.X+w.Width])
|
copy(dst[y*w.Width:(y+1)*w.Width], r.r.data[(w.Y+y)*r.r.grid.Width+w.X:(w.Y+y)*r.r.grid.Width+w.X+w.Width])
|
||||||
}
|
}
|
||||||
@@ -61,13 +96,13 @@ func (f *fakeIO) Open(_ context.Context, path string) (raster.Reader, error) {
|
|||||||
if !ok {
|
if !ok {
|
||||||
return nil, fmt.Errorf("fake: %s not found", path)
|
return nil, fmt.Errorf("fake: %s not found", path)
|
||||||
}
|
}
|
||||||
return &fakeReader{r: r}, nil
|
return &fakeReader{f: f, path: path, r: r}, nil
|
||||||
}
|
}
|
||||||
|
|
||||||
func (f *fakeIO) Create(_ context.Context, path string, g raster.Grid, _ raster.WriterOptions) (raster.Writer, error) {
|
func (f *fakeIO) Create(_ context.Context, path string, g raster.Grid, o raster.WriterOptions) (raster.Writer, error) {
|
||||||
f.mu.Lock()
|
f.mu.Lock()
|
||||||
defer f.mu.Unlock()
|
defer f.mu.Unlock()
|
||||||
r := &fakeRaster{grid: g, data: make([]float64, g.Width*g.Height)}
|
r := &fakeRaster{grid: g, data: make([]float64, g.Width*g.Height), opts: o}
|
||||||
f.files[path] = r
|
f.files[path] = r
|
||||||
return &fakeWriter{f: f, r: r}, nil
|
return &fakeWriter{f: f, r: r}, nil
|
||||||
}
|
}
|
||||||
|
|||||||
Reference in New Issue
Block a user