diff --git a/.env.example b/.env.example index 3765d7d..f0c1f38 100644 --- a/.env.example +++ b/.env.example @@ -9,4 +9,8 @@ S3_SECRET_KEY: minioadmin S3_BUCKET: geofiles S3_USE_SSL: "false" -DB_PORT=5423 \ No newline at end of file +DB_PORT=5423 + +# Landscape Program parallelism (bounds worker peak memory) +LP_MAX_WORKERS=8 +LP_MAX_NODES=2 diff --git a/README.md b/README.md index b37aefb..bfd10e5 100644 --- a/README.md +++ b/README.md @@ -232,6 +232,31 @@ one, an `sld_url` (`/dataset-groups/{id}/sld`). A node failure fails its stage; later stages are not started and the run ends `failed` with the message. Jobs are nacked without requeue, so there is no automatic retry — start a new run. +### Memory and parallelism + +A run has two nested levels of parallelism: nodes of a stage run concurrently, +and each node processes its windows with its own worker pool. Every worker holds +one window buffer per input **and** per output, so the worker's peak resident set +is roughly + +``` +nodes × workers-per-node × window_size² × (inputs + outputs) × 8 bytes +``` + +(the widest node, `20 c_conflicts`, has 6 inputs and 5 outputs). Two environment +variables bound that product, independently of how many cores the host has: + +| Variable | Default | Meaning | +| --- | --- | --- | +| `LP_MAX_WORKERS` | `8` | Total window-worker budget shared by all nodes of a run. | +| `LP_MAX_NODES` | `2` | How many graph nodes of one stage run at the same time. | + +Each node gets `max(1, LP_MAX_WORKERS / LP_MAX_NODES)` workers, further capped by +`GOMAXPROCS` and by the number of windows, so the total stays at +`LP_MAX_WORKERS`. At the defaults and `window_size: 2048` that is about +8 × 2048² × 11 × 8 B ≈ 2.9 GiB worst case — lower `LP_MAX_WORKERS` (or +`window_size`) on a small host, raise it on a large one. + ### Build requirements The raster math uses [godal](https://github.com/airbusgeo/godal), a **cgo** @@ -244,7 +269,7 @@ go build ./... # cgo is on by default for a native build ``` `build/package/Dockerfile` already handles the image: the builder stage adds -`gdal-dev gcc musl-dev pkgconfig` and builds with `CGO_ENABLED=1`, and the +`gdal-dev gcc g++ musl-dev pkgconfig` and builds with `CGO_ENABLED=1`, and the runtime stage installs `gdal` (the shared library godal needs) plus `gdal-tools` (`gdal_translate`/`gdalinfo` for the COG conversion). diff --git a/configs/.env.example b/configs/.env.example index 238bcb4..9a4f79b 100644 --- a/configs/.env.example +++ b/configs/.env.example @@ -28,3 +28,7 @@ RABBITMQ_QUEUE=gis.events.example # Host port mapping for the postgres container (docker-compose) DB_PORT=5432 + +# Landscape Program parallelism (bounds worker peak memory) +LP_MAX_WORKERS=8 +LP_MAX_NODES=2 diff --git a/internal/app/app.go b/internal/app/app.go index 1965f04..3c8ea53 100644 --- a/internal/app/app.go +++ b/internal/app/app.go @@ -12,7 +12,6 @@ import ( "gis/api" "gis/internal/config" "gis/internal/domain" - "gis/internal/landscape" "gis/internal/messaging/rabbitmq" "gis/internal/parser" "gis/internal/platform/logger" @@ -92,9 +91,15 @@ func New(ctx context.Context) (*App, error) { categories: service.NewCategoryService(categoryRepo), datasets: datasetsSvc, groups: service.NewDatasetGroupService(groupRepo, datasetRepo, store), + // A nil ComputeFunc selects landscape.ComputeWithWorkers, so each node + // gets the per-node share of the configured worker budget. landscape: service.NewLandscapeService( lpRepo, datasetRepo, groupRepo, categoryRepo, datasetsSvc, store, - lpJobs, raster.NewGodalIO(), landscape.Compute, log, + lpJobs, raster.NewGodalIO(), nil, log, + service.LandscapeOptions{ + MaxWorkers: cfg.Landscape.MaxWorkers, + MaxConcurrentNodes: cfg.Landscape.MaxConcurrentNodes, + }, ), eventRepo: eventRepo, }, nil diff --git a/internal/config/config.go b/internal/config/config.go index fea10e3..a502098 100644 --- a/internal/config/config.go +++ b/internal/config/config.go @@ -12,10 +12,11 @@ import ( // Config holds all configuration for the application. Values are read from // environment variables; a local .env file (if present) is loaded first. type Config struct { - HTTP HTTPConfig - DB DBConfig - S3 S3Config - RabbitMQ RabbitMQConfig + HTTP HTTPConfig + DB DBConfig + S3 S3Config + RabbitMQ RabbitMQConfig + Landscape LandscapeConfig } // HTTPConfig configures the HTTP server. @@ -58,6 +59,19 @@ type RabbitMQConfig struct { Queue string `env:"RABBITMQ_QUEUE" envDefault:"gis.events.example"` } +// LandscapeConfig bounds the Landscape Program's parallelism, and with it the +// worker's peak memory: a run holds up to MaxConcurrentNodes nodes in flight, +// each running a pool of MaxWorkers/MaxConcurrentNodes goroutines that own one +// window buffer per input and output. The defaults (8 total workers spread +// over 2 nodes) keep that product bounded regardless of how many cores the +// host has. +type LandscapeConfig struct { + // MaxWorkers is the total window-worker budget shared by all nodes of a run. + MaxWorkers int `env:"LP_MAX_WORKERS" envDefault:"8"` + // MaxConcurrentNodes is how many graph nodes of one stage run at a time. + MaxConcurrentNodes int `env:"LP_MAX_NODES" envDefault:"2"` +} + // Load reads configuration from the environment, loading an optional .env file // from the current working directory first. func Load() (*Config, error) { diff --git a/internal/landscape/compute.go b/internal/landscape/compute.go index 4697fe0..d7519cc 100644 --- a/internal/landscape/compute.go +++ b/internal/landscape/compute.go @@ -13,13 +13,27 @@ import ( "golang.org/x/sync/errgroup" ) -// 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 +// MaxWorkersPerNode is the per-node worker cap Compute uses. Peak memory is +// roughly nodes x workers x window^2 x (inputs+outputs), so callers that run +// several nodes at once should divide their total worker budget between them +// and call ComputeWithWorkers instead. +const MaxWorkersPerNode = 8 + +// Compute runs one node over its input rasters with the default per-node +// worker cap (MaxWorkersPerNode). See ComputeWithWorkers. +func Compute(ctx context.Context, rio raster.IO, node Node, inputs map[string]string, outDir string, p Params) (map[string]string, error) { + return ComputeWithWorkers(ctx, rio, node, inputs, outDir, p, MaxWorkersPerNode) +} + +// ComputeWithWorkers runs one node over its input rasters with an explicit +// worker cap. See the package doc for the concurrency model: windows are +// processed by a worker pool of at most workers goroutines; 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. 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) { +// error, no output files are left behind: ComputeWithWorkers best-effort +// removes every output it created before returning. A workers value below 1 is +// treated as 1. +func ComputeWithWorkers(ctx context.Context, rio raster.IO, node Node, inputs map[string]string, outDir string, p Params, workers int) (map[string]string, error) { p = p.WithDefaults() if p.WindowSize <= 0 { // WithDefaults only fills in a zero value; guard against a negative @@ -75,8 +89,8 @@ func Compute(ctx context.Context, rio raster.IO, node Node, inputs map[string]st // 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) + // nodes. The caller's budget is the third bound. + workers = min(runtime.GOMAXPROCS(0), len(windows), max(1, workers)) for w := 0; w < workers; w++ { eg.Go(func() error { readers, err := openAll(gctx, rio, node, paths) diff --git a/internal/landscape/compute_test.go b/internal/landscape/compute_test.go index 5754366..a5a34f1 100644 --- a/internal/landscape/compute_test.go +++ b/internal/landscape/compute_test.go @@ -217,3 +217,36 @@ func TestRemoveOutputs_DeletesFiles(t *testing.T) { } } } + +// TestComputeWithWorkers_RespectsCap checks the caller's worker budget bounds +// the pool: each pool worker opens one reader per input and holds a buffer per +// input and output, so the open count bounds peak memory. checkGrids opens one +// extra set of handles before the pool starts, hence the (workers+1) bound. +func TestComputeWithWorkers_RespectsCap(t *testing.T) { + node, _ := NodeByCode("10") + inputs := len(node.Inputs) + + for _, workers := range []int{1, 2, 4, 0, -3} { + t.Run(fmt.Sprintf("workers=%d", workers), func(t *testing.T) { + fio := newFakeIO() + g := grid(64, 64) + fio.put("z.tif", g, constRaster(64, 64, 2)) + fio.put("u.tif", g, constRaster(64, 64, 4)) + p := DefaultParams() + p.WindowSize = 8 // 64 windows, so the pool is never window-bound + + if _, err := ComputeWithWorkers(context.Background(), fio, node, + map[string]string{OutZIntegrated: "z.tif", OutUIntegrated: "u.tif"}, "/out", p, workers); err != nil { + t.Fatal(err) + } + + effective := max(1, workers) + if got, limit := fio.openCount(), (effective+1)*inputs; got > limit { + t.Fatalf("opened %d readers with workers=%d, want at most %d", got, workers, limit) + } + if fio.openCount() < inputs { + t.Fatalf("opened %d readers, want at least one set of %d", fio.openCount(), inputs) + } + }) + } +} diff --git a/internal/landscape/fakeio_test.go b/internal/landscape/fakeio_test.go index 61e1ce8..b4f4b94 100644 --- a/internal/landscape/fakeio_test.go +++ b/internal/landscape/fakeio_test.go @@ -19,6 +19,17 @@ type fakeIO struct { // mid-run I/O failure. failReadAfter map[string]int readCounts map[string]int + + // opens counts every successful Open. Each pool worker opens one reader + // per input, so the count bounds the number of concurrent handles. + opens int +} + +// openCount returns how many readers have been opened so far. +func (f *fakeIO) openCount() int { + f.mu.Lock() + defer f.mu.Unlock() + return f.opens } type fakeRaster struct { @@ -96,6 +107,7 @@ func (f *fakeIO) Open(_ context.Context, path string) (raster.Reader, error) { if !ok { return nil, fmt.Errorf("fake: %s not found", path) } + f.opens++ return &fakeReader{f: f, path: path, r: r}, nil } diff --git a/internal/landscape/kernels.go b/internal/landscape/kernels.go index 537ec32..833b83d 100644 --- a/internal/landscape/kernels.go +++ b/internal/landscape/kernels.go @@ -34,7 +34,9 @@ func kernelZIntegrated(in [][]float64, out [][]int16, _ Params) { out[0][i] = NoData continue } - score := float32(0.25)*float32(z1) + float32(0.25)*float32(z2) + float32(0.25)*float32(z3) + float32(0.25)*float32(z4) + // explicit conversions forbid FMA fusion so results match numpy's per-op rounding + score := float32(float32(0.25)*float32(z1)) + float32(float32(0.25)*float32(z2)) + + float32(float32(0.25)*float32(z3)) + float32(float32(0.25)*float32(z4)) out[0][i] = classify(score) } } @@ -49,14 +51,15 @@ func kernelUIntegrated(in [][]float64, out [][]int16, _ Params) { out[0][i] = NoData continue } - num := float32(0.25)*float32(water) + float32(0.25)*float32(soil) + // explicit conversions forbid FMA fusion so results match numpy's per-op rounding + num := float32(float32(0.25)*float32(water)) + float32(float32(0.25)*float32(soil)) den := float32(0.5) if cropOK { - num += float32(0.25) * float32(crop) + num += float32(float32(0.25) * float32(crop)) den += 0.25 } if vegOK { - num += float32(0.25) * float32(veg) + num += float32(float32(0.25) * float32(veg)) den += 0.25 } out[0][i] = classify(num / den) diff --git a/internal/messaging/rabbitmq/landscape_consumer.go b/internal/messaging/rabbitmq/landscape_consumer.go index bf3c5d9..67a09ad 100644 --- a/internal/messaging/rabbitmq/landscape_consumer.go +++ b/internal/messaging/rabbitmq/landscape_consumer.go @@ -5,6 +5,7 @@ import ( "encoding/json" "fmt" "log/slog" + "time" "github.com/google/uuid" amqp "github.com/rabbitmq/amqp091-go" @@ -15,6 +16,9 @@ const ( LandscapeRunRoutingKey = "landscape.run" // LandscapeRunQueue receives Landscape Program run jobs. LandscapeRunQueue = "gis.landscape.run" + // markRunFailedTimeout bounds the best-effort terminal status write made + // after Execute failed, on a context detached from the delivery's. + markRunFailedTimeout = 10 * time.Second ) // LandscapeJob is the message body for a Landscape Program run job. @@ -72,9 +76,14 @@ func (h landscapeHandler) Handle(ctx context.Context, d amqp.Delivery) error { // status write), since Execute already marks the run failed for every // node failure. Best-effort: log if marking fails, but still return the // original error so the delivery is nacked. - if markErr := h.proc.MarkRunFailed(ctx, job.RunID, err.Error()); markErr != nil { + // Detached from ctx: Execute may well have failed *because* ctx was + // cancelled (shutdown), and the write must still land or the run stays + // "running" forever. + mctx, cancel := context.WithTimeout(context.WithoutCancel(ctx), markRunFailedTimeout) + if markErr := h.proc.MarkRunFailed(mctx, job.RunID, err.Error()); markErr != nil { h.log.Error("mark landscape run failed", "run_id", job.RunID, "error", markErr) } + cancel() } return err } diff --git a/internal/service/landscape.go b/internal/service/landscape.go index 6d2d66b..e6e1810 100644 --- a/internal/service/landscape.go +++ b/internal/service/landscape.go @@ -11,7 +11,6 @@ import ( "log/slog" "os" "path/filepath" - "runtime" "sort" "sync" "time" @@ -29,6 +28,47 @@ import ( // mirrors landscape.Compute so tests can substitute a stub. type ComputeFunc func(ctx context.Context, rio raster.IO, node landscape.Node, inputs map[string]string, outDir string, p landscape.Params) (map[string]string, error) +// ComputeWithWorkersFunc is ComputeFunc plus an explicit per-node worker cap. +// It mirrors landscape.ComputeWithWorkers. +type ComputeWithWorkersFunc func(ctx context.Context, rio raster.IO, node landscape.Node, inputs map[string]string, outDir string, p landscape.Params, workers int) (map[string]string, error) + +// Default bounds on a run's parallelism, mirroring config.LandscapeConfig. +const ( + // DefaultMaxWorkers is the total window-worker budget of one run. + DefaultMaxWorkers = 8 + // DefaultMaxConcurrentNodes is how many nodes of a stage run at a time. + DefaultMaxConcurrentNodes = 2 +) + +// terminalWriteTimeout bounds a terminal status write made on a context that +// has already been cancelled (see terminalCtx). +const terminalWriteTimeout = 10 * time.Second + +// LandscapeOptions bounds a run's parallelism, and with it the worker's peak +// memory: MaxConcurrentNodes nodes run at a time, and they share the +// MaxWorkers window-worker budget. Zero or negative values fall back to the +// defaults. +type LandscapeOptions struct { + MaxWorkers int + MaxConcurrentNodes int +} + +func (o LandscapeOptions) withDefaults() LandscapeOptions { + if o.MaxWorkers < 1 { + o.MaxWorkers = DefaultMaxWorkers + } + if o.MaxConcurrentNodes < 1 { + o.MaxConcurrentNodes = DefaultMaxConcurrentNodes + } + return o +} + +// workersPerNode splits the total worker budget over the concurrently running +// nodes, so nodes x workers stays at MaxWorkers. Always at least 1. +func (o LandscapeOptions) workersPerNode() int { + return max(1, o.MaxWorkers/o.MaxConcurrentNodes) +} + // LandscapeRepository is the persistence behaviour LandscapeService needs. type LandscapeRepository interface { CreateRun(ctx context.Context, run domain.LPRun) (domain.LPRun, error) @@ -115,15 +155,24 @@ type LandscapeService struct { store ObjectStore jobs LandscapeJobEnqueuer rio raster.IO - compute ComputeFunc - log *slog.Logger + // compute is the legacy worker-count-agnostic entry point, kept so an + // existing ComputeFunc still wires up; runNode always calls + // computeWithWorkers, which wraps compute when one was provided. + compute ComputeFunc + computeWithWorkers ComputeWithWorkersFunc + opts LandscapeOptions + log *slog.Logger } // NewLandscapeService wires the run repository, the dataset and group // repositories (input resolution), the category reader (the "landscape" // category every result is filed under), the dataset publisher (result // datasets and slot uploads), the object store (input downloads), the job -// enqueuer, and the raster IO / compute function the graph runs on. +// enqueuer, and the raster IO / compute function the graph runs on. Pass a nil +// compute (the production wiring) to run landscape.ComputeWithWorkers with the +// per-node share of opts.MaxWorkers; a non-nil ComputeFunc is honoured for +// backward compatibility but never sees the worker cap. opts bounds the run's +// parallelism; zero fields fall back to the package defaults. func NewLandscapeService( runs LandscapeRepository, datasets DatasetRepository, @@ -135,18 +184,29 @@ func NewLandscapeService( rio raster.IO, compute ComputeFunc, log *slog.Logger, + opts LandscapeOptions, ) *LandscapeService { + // Default to the workers-aware entry point; wrap a supplied legacy + // ComputeFunc (which has no worker parameter) when one is given. + withWorkers := ComputeWithWorkersFunc(landscape.ComputeWithWorkers) + if compute != nil { + withWorkers = func(ctx context.Context, rio raster.IO, node landscape.Node, inputs map[string]string, outDir string, p landscape.Params, _ int) (map[string]string, error) { + return compute(ctx, rio, node, inputs, outDir, p) + } + } return &LandscapeService{ - runs: runs, - datasets: datasets, - groups: groups, - categories: categories, - publisher: publisher, - store: store, - jobs: jobs, - rio: rio, - compute: compute, - log: log, + runs: runs, + datasets: datasets, + groups: groups, + categories: categories, + publisher: publisher, + store: store, + jobs: jobs, + rio: rio, + compute: compute, + computeWithWorkers: withWorkers, + opts: opts.withDefaults(), + log: log, } } @@ -344,7 +404,9 @@ func (s *LandscapeService) Results(ctx context.Context) ([]ResultInfo, error) { // MarkRunFailed records a terminal failure on a run so it does not stay stuck // in an in-progress status. func (s *LandscapeService) MarkRunFailed(ctx context.Context, runID uuid.UUID, reason string) error { - return s.runs.SetRunStatus(ctx, runID, domain.LPRunFailed, &reason) + wctx, cancel := terminalCtx(ctx) + defer cancel() + return s.runs.SetRunStatus(wctx, runID, domain.LPRunFailed, &reason) } // Execute runs the whole graph for one run: stages in order, nodes within a @@ -355,6 +417,13 @@ func (s *LandscapeService) Execute(ctx context.Context, runID uuid.UUID) error { if err != nil { return err } + // A redelivered or stale message must not rerun a finished run. "running" + // is allowed through so a run whose worker crashed mid-flight still + // recovers on redelivery. + if run.Status != domain.LPRunQueued && run.Status != domain.LPRunRunning { + s.log.Info("landscape run already terminal, ignoring redelivery", "run_id", runID, "status", run.Status) + return nil + } var params landscape.Params if len(run.Params) > 0 { @@ -382,15 +451,18 @@ func (s *LandscapeService) Execute(ctx context.Context, runID uuid.UUID) error { state := &runState{produced: map[string]uuid.UUID{}} for _, stage := range landscape.Stages() { - var g errgroup.Group - // Each node's Compute spins its own worker pool (up to - // min(GOMAXPROCS, windows, 8) workers, each holding a window buffer per - // input band), so peak memory is the product of the two limits. Bound - // the node fan-out here so that product stays sane: without a limit all - // four stage-0 nodes would run at once and multiply Compute's own cap. - g.SetLimit(maxConcurrentNodes()) + // WithContext so the first failing node cancels its siblings: Compute + // honours cancellation, so a doomed stage stops reading rasters + // instead of running to completion. Terminal status writes use + // terminalCtx, so they still land on the cancelled context. + g, gctx := errgroup.WithContext(ctx) + // Each node's Compute spins its own worker pool (each worker holding a + // window buffer per input band), so peak memory is the product of the + // two limits. Bound the node fan-out here and split the worker budget + // over it so that product stays at opts.MaxWorkers. + g.SetLimit(s.opts.MaxConcurrentNodes) for _, node := range stage { - g.Go(func() error { return s.runNode(ctx, run, node, params, cat, state) }) + g.Go(func() error { return s.runNode(gctx, run, node, params, cat, state) }) } // Wait for every node of the stage: each one records its own outcome, // and a sibling's failure must not leave a node half-recorded. @@ -402,10 +474,20 @@ func (s *LandscapeService) Execute(ctx context.Context, runID uuid.UUID) error { return s.runs.SetRunStatus(ctx, runID, domain.LPRunSucceeded, nil) } +// terminalCtx derives a bounded context that outlives the cancellation of ctx. +// Terminal status writes must happen even when the run failed *because* the +// context was cancelled (shutdown, a sibling node failing the stage); on the +// original context they would be refused and leave the run stuck in "running". +func terminalCtx(ctx context.Context) (context.Context, context.CancelFunc) { + return context.WithTimeout(context.WithoutCancel(ctx), terminalWriteTimeout) +} + // failRun marks the run failed with cause's message and returns cause. func (s *LandscapeService) failRun(ctx context.Context, runID uuid.UUID, cause error) error { msg := cause.Error() - if err := s.runs.SetRunStatus(ctx, runID, domain.LPRunFailed, &msg); err != nil { + wctx, cancel := terminalCtx(ctx) + defer cancel() + if err := s.runs.SetRunStatus(wctx, runID, domain.LPRunFailed, &msg); err != nil { s.log.Error("mark landscape run failed", "run_id", runID, "error", err) } return cause @@ -473,7 +555,9 @@ func (s *LandscapeService) runNode(ctx context.Context, run domain.LPRun, node l // finishNodeFailed records the failure on the node and returns the cause. func (s *LandscapeService) finishNodeFailed(ctx context.Context, runID uuid.UUID, node landscape.Node, cause error) error { msg := cause.Error() - if err := s.runs.SetNodeFinished(ctx, runID, node.Code, domain.LPNodeFailed, nil, &msg); err != nil { + wctx, cancel := terminalCtx(ctx) + defer cancel() + if err := s.runs.SetNodeFinished(wctx, runID, node.Code, domain.LPNodeFailed, nil, &msg); err != nil { s.log.Error("mark landscape node failed", "run_id", runID, "node", node.Code, "error", err) } return cause @@ -583,7 +667,7 @@ func (s *LandscapeService) computeNode( paths[group] = path } - outPaths, err := s.compute(ctx, s.rio, node, paths, dir, params) + outPaths, err := s.computeWithWorkers(ctx, s.rio, node, paths, dir, params, s.opts.workersPerNode()) if err != nil { return nil, err } @@ -693,21 +777,6 @@ func knownGroup(code string) bool { return ok } -// maxConcurrentNodes caps how many graph nodes of one stage run at the same -// time. Compute already parallelises within a node, so half the CPUs (at most -// 4, the widest stage) is enough to keep the machine busy without letting the -// two levels of parallelism multiply into an out-of-memory. -func maxConcurrentNodes() int { - n := runtime.GOMAXPROCS(0) / 2 - if n < 1 { - n = 1 - } - if n > 4 { - n = 4 - } - return n -} - // maxStage returns the highest stage number in the graph. func maxStage() int { highest := 0 diff --git a/internal/service/landscape_test.go b/internal/service/landscape_test.go index 94d34a0..1fc4244 100644 --- a/internal/service/landscape_test.go +++ b/internal/service/landscape_test.go @@ -65,12 +65,63 @@ func (r *lockedDatasetRepo) LatestInGroup(ctx context.Context, groupID uuid.UUID } // stubLPRepo is an in-memory LandscapeRepository for tests. +// lpStatusWrite records one SetRunStatus call and whether the context it was +// handed was still live. +type lpStatusWrite struct { + status domain.LPRunStatus + ctxErr error +} + +// lpNodeWrite records one SetNodeFinished call the same way. +type lpNodeWrite struct { + code string + status domain.LPNodeStatus + ctxErr error +} + type stubLPRepo struct { mu sync.Mutex runs map[uuid.UUID]*domain.LPRun order []uuid.UUID reuseLookups []string clock int64 + statusWrites []lpStatusWrite + nodeWrites []lpNodeWrite +} + +// statusesWritten returns the run statuses written so far, in order. +func (r *stubLPRepo) statusesWritten() []domain.LPRunStatus { + r.mu.Lock() + defer r.mu.Unlock() + out := make([]domain.LPRunStatus, 0, len(r.statusWrites)) + for _, w := range r.statusWrites { + out = append(out, w.status) + } + return out +} + +// statusWrite returns the last write recording the given status. +func (r *stubLPRepo) statusWrite(status domain.LPRunStatus) (lpStatusWrite, bool) { + r.mu.Lock() + defer r.mu.Unlock() + for i := len(r.statusWrites) - 1; i >= 0; i-- { + if r.statusWrites[i].status == status { + return r.statusWrites[i], true + } + } + return lpStatusWrite{}, false +} + +// nodeWrite returns the last SetNodeFinished write recording the given status. +func (r *stubLPRepo) nodeWrite(status domain.LPNodeStatus) (lpNodeWrite, bool) { + r.mu.Lock() + defer r.mu.Unlock() + for i := len(r.nodeWrites) - 1; i >= 0; i-- { + if r.nodeWrites[i].status == status { + return r.nodeWrites[i], true + } + } + return lpNodeWrite{}, false } func newStubLPRepo() *stubLPRepo { @@ -170,9 +221,10 @@ func (r *stubLPRepo) CountRuns(_ context.Context) (int, error) { return len(r.runs), nil } -func (r *stubLPRepo) SetRunStatus(_ context.Context, id uuid.UUID, status domain.LPRunStatus, errMsg *string) error { +func (r *stubLPRepo) SetRunStatus(ctx context.Context, id uuid.UUID, status domain.LPRunStatus, errMsg *string) error { r.mu.Lock() defer r.mu.Unlock() + r.statusWrites = append(r.statusWrites, lpStatusWrite{status: status, ctxErr: ctx.Err()}) run, ok := r.runs[id] if !ok { return domain.ErrNotFound @@ -219,9 +271,10 @@ func (r *stubLPRepo) SetNodeStarted(_ context.Context, runID uuid.UUID, code, fi return nil } -func (r *stubLPRepo) SetNodeFinished(_ context.Context, runID uuid.UUID, code string, status domain.LPNodeStatus, outputs []domain.LPDatasetRef, errMsg *string) error { +func (r *stubLPRepo) SetNodeFinished(ctx context.Context, runID uuid.UUID, code string, status domain.LPNodeStatus, outputs []domain.LPDatasetRef, errMsg *string) error { r.mu.Lock() defer r.mu.Unlock() + r.nodeWrites = append(r.nodeWrites, lpNodeWrite{code: code, status: status, ctxErr: ctx.Err()}) n, err := r.node(runID, code) if err != nil { return err @@ -354,8 +407,11 @@ func (j *stubLPJobs) EnqueueRun(_ context.Context, runID uuid.UUID) error { // stubComputeCall records one invocation of the compute stub. type stubComputeCall struct { - node string - inputs map[string]string + node string + // workers is the per-node worker cap the service passed, or -1 when the + // call came through the legacy worker-agnostic ComputeFunc. + workers int + inputs map[string]string // bodies holds each input file's content at call time; seedDataset stores // each object under its own storage key as its body, so a test can tell // which dataset a node actually downloaded. @@ -369,11 +425,47 @@ type stubCompute struct { mu sync.Mutex calls []stubComputeCall errs map[string]error + // waitForCancel names nodes whose compute blocks until its context is + // cancelled (or the deadline below elapses), so a test can observe whether + // a sibling's failure cancelled them. + waitForCancel map[string]bool + sawCancel map[string]bool } -func newStubCompute() *stubCompute { return &stubCompute{errs: map[string]error{}} } +func newStubCompute() *stubCompute { + return &stubCompute{errs: map[string]error{}, waitForCancel: map[string]bool{}, sawCancel: map[string]bool{}} +} + +// fn is a legacy service.ComputeFunc: it sees no worker cap. +func (c *stubCompute) fn(ctx context.Context, rio raster.IO, node landscape.Node, inputs map[string]string, outDir string, p landscape.Params) (map[string]string, error) { + return c.fnWithWorkers(ctx, rio, node, inputs, outDir, p, -1) +} + +// observedCancel reports whether the node's compute saw a cancelled context. +func (c *stubCompute) observedCancel(code string) bool { + c.mu.Lock() + defer c.mu.Unlock() + return c.sawCancel[code] +} + +func (c *stubCompute) fnWithWorkers(ctx context.Context, _ raster.IO, node landscape.Node, inputs map[string]string, outDir string, p landscape.Params, workers int) (map[string]string, error) { + c.mu.Lock() + wait := c.waitForCancel[node.Code] + c.mu.Unlock() + if wait { + select { + case <-ctx.Done(): + case <-time.After(5 * time.Second): + } + c.mu.Lock() + c.sawCancel[node.Code] = ctx.Err() != nil + c.mu.Unlock() + } + // The real Compute honours cancellation, so the stub must too. + if err := ctx.Err(); err != nil { + return nil, err + } -func (c *stubCompute) fn(_ context.Context, _ raster.IO, node landscape.Node, inputs map[string]string, outDir string, p landscape.Params) (map[string]string, error) { in := map[string]string{} bodies := map[string]string{} for k, v := range inputs { @@ -386,7 +478,7 @@ func (c *stubCompute) fn(_ context.Context, _ raster.IO, node landscape.Node, in } c.mu.Lock() - c.calls = append(c.calls, stubComputeCall{node: node.Code, inputs: in, bodies: bodies, params: p}) + c.calls = append(c.calls, stubComputeCall{node: node.Code, workers: workers, inputs: in, bodies: bodies, params: p}) err := c.errs[node.Code] c.mu.Unlock() if err != nil { @@ -453,6 +545,12 @@ type lpFixture struct { } func newLPFixture(t *testing.T) *lpFixture { + t.Helper() + return newLPFixtureWith(t, LandscapeOptions{}) +} + +// newLPFixtureWith builds the fixture with explicit parallelism options. +func newLPFixtureWith(t *testing.T, opts LandscapeOptions) *lpFixture { t.Helper() ctx := context.Background() @@ -482,7 +580,11 @@ func newLPFixture(t *testing.T) *lpFixture { comp := newStubCompute() log := slog.New(slog.NewTextHandler(io.Discard, nil)) - svc := NewLandscapeService(runs, datasets, groups, cats, pub, store, jobs, nil, comp.fn, log) + svc := NewLandscapeService(runs, datasets, groups, cats, pub, store, jobs, nil, nil, log, opts) + // Substitute the workers-aware stub so tests can see the per-node cap the + // service computes; NewLandscapeService's legacy-ComputeFunc wrapping is + // covered separately. + svc.computeWithWorkers = comp.fnWithWorkers return &lpFixture{ svc: svc, runs: runs, datasets: datasets, groups: groups, cats: cats, pub: pub, store: store, jobs: jobs, compute: comp, category: cat, @@ -1216,3 +1318,203 @@ func TestLandscapeService_Execute_OverridePinsInputOverFresherDataset(t *testing t.Fatalf("node 01 input for %s = %v, want the freshest %v", landscape.SlotUVeg, got, newest.ID) } } + +// F2: a redelivered message for a finished run must be acked without rerunning. +func TestLandscapeService_Execute_IgnoresTerminalRun(t *testing.T) { + f := newLPFixture(t) + ctx := context.Background() + f.seedAllSlots(t) + + run, err := f.svc.CreateRun(ctx, RunInput{}) + if err != nil { + t.Fatalf("CreateRun: %v", err) + } + if err := f.runs.SetRunStatus(ctx, run.ID, domain.LPRunSucceeded, nil); err != nil { + t.Fatal(err) + } + f.runs.mu.Lock() + f.runs.statusWrites = nil + f.runs.mu.Unlock() + + if err := f.svc.Execute(ctx, run.ID); err != nil { + t.Fatalf("Execute on a terminal run must ack, got %v", err) + } + if n := f.compute.count(); n != 0 { + t.Fatalf("compute called %d times for a terminal run", n) + } + if got := f.runs.statusesWritten(); len(got) != 0 { + t.Fatalf("no status write expected, got %v", got) + } + got, err := f.svc.GetRun(ctx, run.ID) + if err != nil { + t.Fatal(err) + } + if got.Status != domain.LPRunSucceeded { + t.Fatalf("run status changed to %s", got.Status) + } +} + +// F2: a "running" run is still executed, so a worker that crashed mid-run +// recovers on redelivery. +func TestLandscapeService_Execute_RerunsRunningRun(t *testing.T) { + f := newLPFixture(t) + ctx := context.Background() + f.seedAllSlots(t) + + run, err := f.svc.CreateRun(ctx, RunInput{}) + if err != nil { + t.Fatalf("CreateRun: %v", err) + } + if err := f.runs.SetRunStatus(ctx, run.ID, domain.LPRunRunning, nil); err != nil { + t.Fatal(err) + } + if err := f.svc.Execute(ctx, run.ID); err != nil { + t.Fatalf("Execute: %v", err) + } + if f.compute.count() == 0 { + t.Fatal("a running run must be executed on redelivery") + } +} + +// F4: terminal status writes must land even when the caller's context is dead. +func TestLandscapeService_Execute_TerminalWritesSurviveCancelledContext(t *testing.T) { + f := newLPFixture(t) + f.seedAllSlots(t) + + run, err := f.svc.CreateRun(context.Background(), RunInput{}) + if err != nil { + t.Fatalf("CreateRun: %v", err) + } + + ctx, cancel := context.WithCancel(context.Background()) + cancel() + if err := f.svc.Execute(ctx, run.ID); err == nil { + t.Fatal("expected Execute to fail on a cancelled context") + } + + got, err := f.svc.GetRun(context.Background(), run.ID) + if err != nil { + t.Fatal(err) + } + if got.Status != domain.LPRunFailed { + t.Fatalf("run status = %s, want failed", got.Status) + } + w, ok := f.runs.statusWrite(domain.LPRunFailed) + if !ok { + t.Fatal("no failed status write recorded") + } + if w.ctxErr != nil { + t.Fatalf("terminal run status write got a cancelled context: %v", w.ctxErr) + } + nw, ok := f.runs.nodeWrite(domain.LPNodeFailed) + if !ok { + t.Fatal("no failed node write recorded") + } + if nw.ctxErr != nil { + t.Fatalf("terminal node write got a cancelled context: %v", nw.ctxErr) + } +} + +// F4: MarkRunFailed is the consumer's last resort and must not need a live ctx. +func TestLandscapeService_MarkRunFailed_IgnoresCancelledContext(t *testing.T) { + f := newLPFixture(t) + run, err := f.svc.CreateRun(context.Background(), RunInput{}) + if err != nil { + t.Fatalf("CreateRun: %v", err) + } + ctx, cancel := context.WithCancel(context.Background()) + cancel() + if err := f.svc.MarkRunFailed(ctx, run.ID, "boom"); err != nil { + t.Fatalf("MarkRunFailed: %v", err) + } + w, ok := f.runs.statusWrite(domain.LPRunFailed) + if !ok { + t.Fatal("no failed status write recorded") + } + if w.ctxErr != nil { + t.Fatalf("MarkRunFailed wrote on a cancelled context: %v", w.ctxErr) + } +} + +// F5: the per-node worker cap handed to Compute is the run's budget split over +// the concurrently running nodes. +func TestLandscapeService_Execute_SplitsWorkerBudgetOverNodes(t *testing.T) { + cases := []struct { + name string + opts LandscapeOptions + want int + }{ + {"defaults", LandscapeOptions{}, DefaultMaxWorkers / DefaultMaxConcurrentNodes}, + {"explicit", LandscapeOptions{MaxWorkers: 12, MaxConcurrentNodes: 3}, 4}, + {"single node keeps the whole budget", LandscapeOptions{MaxWorkers: 6, MaxConcurrentNodes: 1}, 6}, + {"never below one", LandscapeOptions{MaxWorkers: 1, MaxConcurrentNodes: 4}, 1}, + } + for _, c := range cases { + t.Run(c.name, func(t *testing.T) { + f := newLPFixtureWith(t, c.opts) + f.seedAllSlots(t) + got := f.runOnce(t, RunInput{}) + if got.Status != domain.LPRunSucceeded { + t.Fatalf("run status = %s: %v", got.Status, got.Error) + } + f.compute.mu.Lock() + defer f.compute.mu.Unlock() + if len(f.compute.calls) == 0 { + t.Fatal("compute was never called") + } + for _, call := range f.compute.calls { + if call.workers != c.want { + t.Fatalf("node %s got workers=%d, want %d", call.node, call.workers, c.want) + } + } + }) + } +} + +// F6: a failing node cancels its siblings instead of letting them run on. +func TestLandscapeService_Execute_FailingNodeCancelsSiblings(t *testing.T) { + // Four stage-0 nodes, all in flight at once; "00" fails immediately while + // "01" blocks waiting for its context. + f := newLPFixtureWith(t, LandscapeOptions{MaxWorkers: 8, MaxConcurrentNodes: 4}) + f.seedAllSlots(t) + f.compute.errs["00"] = errors.New("boom") + f.compute.waitForCancel["01"] = true + + got := f.runOnce(t, RunInput{}) + if got.Status != domain.LPRunFailed { + t.Fatalf("run status = %s, want failed", got.Status) + } + if !f.compute.observedCancel("01") { + t.Fatal("sibling node 01 did not observe the cancellation of its context") + } + for _, n := range got.Nodes { + if n.NodeCode == "01" && n.Status == domain.LPNodeSucceeded { + t.Fatal("cancelled sibling must not be recorded as succeeded") + } + } +} + +// F5: a legacy ComputeFunc still wires up through NewLandscapeService. +func TestNewLandscapeService_WrapsLegacyComputeFunc(t *testing.T) { + f := newLPFixture(t) + f.seedAllSlots(t) + // Rebuild the service the production-adjacent way: a legacy ComputeFunc, + // no direct field substitution. + log := slog.New(slog.NewTextHandler(io.Discard, nil)) + svc := NewLandscapeService(f.runs, f.datasets, f.groups, f.cats, f.pub, f.store, f.jobs, nil, f.compute.fn, log, LandscapeOptions{}) + + run, err := svc.CreateRun(context.Background(), RunInput{}) + if err != nil { + t.Fatalf("CreateRun: %v", err) + } + if err := svc.Execute(context.Background(), run.ID); err != nil { + t.Fatalf("Execute: %v", err) + } + call, ok := f.compute.callFor("00") + if !ok { + t.Fatal("legacy ComputeFunc was never called") + } + if call.workers != -1 { + t.Fatalf("legacy ComputeFunc should not see a worker cap, got %d", call.workers) + } +}