From 642777f1d240976262fc0fb283922cbbf2367e1a Mon Sep 17 00:00:00 2001 From: Edoardo Vacchi Date: Wed, 1 Jul 2026 10:20:38 +0200 Subject: [PATCH 1/2] refactor: introduce result collector Signed-off-by: Edoardo Vacchi --- internal/processor/worker/collector.go | 162 ++++++++++ internal/processor/worker/executor.go | 357 ++++++++------------- internal/processor/worker/executor_test.go | 162 ++++------ 3 files changed, 364 insertions(+), 317 deletions(-) create mode 100644 internal/processor/worker/collector.go diff --git a/internal/processor/worker/collector.go b/internal/processor/worker/collector.go new file mode 100644 index 000000000..68a7e9d47 --- /dev/null +++ b/internal/processor/worker/collector.go @@ -0,0 +1,162 @@ +/* +Copyright 2026 The llm-d Authors + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +*/ + +package worker + +import ( + "bufio" + "context" + "encoding/json" + "sync" + + "github.com/go-logr/logr" + + batch_types "github.com/llm-d/llm-d-batch-gateway/internal/shared/types" +) + +// ResultItem is the outcome of a single inference request. +// Produced by processModel goroutines, consumed by resultCollector. +type ResultItem struct { + RequestID string + CustomID string + Response *batch_types.ResponseData + Error *OutputError + HadCapacityRetry bool + ModelID string +} + +// OutputError is the error structure written to JSONL output/error files. +type OutputError struct { + Code string `json:"code"` + Message string `json:"message"` +} + +func (r *ResultItem) isSuccess() bool { + return r.Error == nil && r.Response != nil && r.Response.StatusCode == 200 +} + +func resultToOutputLine(r *ResultItem) *outputLine { + var outErr *outputError + if r.Error != nil { + outErr = &outputError{Code: r.Error.Code, Message: r.Error.Message} + } + return &outputLine{ + ID: r.RequestID, + CustomID: r.CustomID, + Response: r.Response, + Error: outErr, + } +} + +// resultCollector reads ResultItems from an internal channel, marshals each to +// JSONL, writes to the appropriate file (output or error), and records progress. +// It runs as a single goroutine. +// +// On write/marshal errors the collector calls abortFn to stop dispatch, logs +// the error, and continues draining so senders don't deadlock. The channel +// buffer (1024) absorbs bursts, but senders will block if it fills. +// +// Usage: +// +// c := newResultCollector(outputBuf, errorBuf, progress, logger, abortFn) +// c.start() +// // ... call c.collect(result) from any goroutine ... +// c.flush() // closes channel, waits for goroutine to finish +type resultCollector struct { + outputWriter *bufio.Writer + errorWriter *bufio.Writer + progress *executionProgress + logger logr.Logger + abortFn context.CancelFunc + abortOnce sync.Once + + ch chan *ResultItem + done chan struct{} +} + +func newResultCollector(outputWriter, errorWriter *bufio.Writer, progress *executionProgress, logger logr.Logger, abortFn context.CancelFunc) *resultCollector { + if abortFn == nil { + panic("resultCollector: abortFn cannot be nil") + } + return &resultCollector{ + outputWriter: outputWriter, + errorWriter: errorWriter, + progress: progress, + logger: logger, + abortFn: abortFn, + ch: make(chan *ResultItem, 1024), + done: make(chan struct{}), + } +} + +// start launches the collector goroutine. The context is used for progress +// updates to the status store. +func (c *resultCollector) start(ctx context.Context) { + go func() { + c.run(ctx) + close(c.done) + }() +} + +// collect sends a result to the collector goroutine. +func (c *resultCollector) collect(result *ResultItem) { + c.ch <- result +} + +// flush closes the result channel and waits for the collector goroutine to +// finish writing all buffered results and flushing the underlying writers. +func (c *resultCollector) flush() { + close(c.ch) + <-c.done +} + +func (c *resultCollector) abort() { + c.abortOnce.Do(c.abortFn) +} + +func (c *resultCollector) run(ctx context.Context) { + for result := range c.ch { + line := resultToOutputLine(result) + + lineBytes, err := json.Marshal(line) + if err != nil { + c.logger.Error(err, "Failed to marshal output line", "customId", result.CustomID) + c.abort() + continue + } + lineBytes = append(lineBytes, '\n') + + isError := line.Error != nil + writer := c.outputWriter + if isError { + writer = c.errorWriter + } + if _, err := writer.Write(lineBytes); err != nil { + c.logger.Error(err, "Failed to write output line", "customId", result.CustomID) + c.abort() + continue + } + + c.progress.record(ctx, result.isSuccess()) + } + + if err := c.outputWriter.Flush(); err != nil { + c.logger.Error(err, "Failed to flush output file (partial results may be truncated)") + } + if err := c.errorWriter.Flush(); err != nil { + c.logger.Error(err, "Failed to flush error file (partial results may be truncated)") + } +} diff --git a/internal/processor/worker/executor.go b/internal/processor/worker/executor.go index 9e2dd9f15..be6fba7f2 100644 --- a/internal/processor/worker/executor.go +++ b/internal/processor/worker/executor.go @@ -43,29 +43,6 @@ import ( "github.com/llm-d/llm-d-batch-gateway/pkg/clients/inference" ) -// outputWriters holds the buffered writers and their mutexes for the output and error JSONL files. -// A single instance is created per job and shared across model goroutines. -type outputWriters struct { - output *bufio.Writer - outputMu sync.Mutex - errors *bufio.Writer - errorsMu sync.Mutex -} - -// write writes line to the error file if isError is true, otherwise to the output file. -func (w *outputWriters) write(line []byte, isError bool) error { - if isError { - w.errorsMu.Lock() - defer w.errorsMu.Unlock() - _, err := w.errors.Write(line) - return err - } - w.outputMu.Lock() - defer w.outputMu.Unlock() - _, err := w.output.Write(line) - return err -} - // outputLine represents a single line in the output JSONL file following the OpenAI batch output format. type outputLine struct { ID string `json:"id"` @@ -234,19 +211,11 @@ func (p *Processor) executeJob(ctx, sloCtx, userCancelCtx, requestAbortCtx conte } defer errorFile.Close() - writers := &outputWriters{ - output: bufio.NewWriterSize(outputFile, 1024*1024), - errors: bufio.NewWriterSize(errorFile, 1024*1024), - } - plansDir, err := p.jobPlansDir(params.jobInfo.JobID, params.jobInfo.TenantID) if err != nil { return nil, err } - // requestAbortCtx and requestAbortFn are set in runJob before watchCancel starts, - // eliminating the race window where a cancel event could arrive before the fn is assigned. - progress := &executionProgress{ total: modelMap.LineCount, updater: params.updater, @@ -255,7 +224,19 @@ func (p *Processor) executeJob(ctx, sloCtx, userCancelCtx, requestAbortCtx conte // Seed with requests already rejected during ingestion (model not found). progress.failed.Store(modelMap.RejectedCount) - errCh := make(chan error, len(modelMap.SafeToModel)) + abortFn := params.requestAbortFn + // This only happens in tests. + if abortFn == nil { + abortFn = func() {} + } + collector := newResultCollector( + bufio.NewWriterSize(outputFile, 1024*1024), + bufio.NewWriterSize(errorFile, 1024*1024), + progress, + logger, + abortFn, + ) + collector.start(ctx) passThroughHeaders := params.jobInfo.PassThroughHeaders if len(passThroughHeaders) > 0 { @@ -266,19 +247,40 @@ func (p *Processor) executeJob(ctx, sloCtx, userCancelCtx, requestAbortCtx conte logger.V(logging.DEBUG).Info("pass-through headers attached to job", "headerNames", headerNames) } - // User-initiated cancellation: watchCancel calls userCancelFn() only. - // context.AfterFunc(userCancelCtx, requestAbortFn) — wired in runJob — then - // cancels requestAbortCtx to stop the dispatch loop. userCancelCtx is isolated - // from sloCtx (derived from context.Background), so SLO expiry and SIGTERM do - // not set it. processModel's drain phase checks sloCtx.Err() vs - // userCancelCtx.Err() to choose the right error code (errExpired vs errCancelled). tenantID := params.jobInfo.TenantID + // Error handler: iterates over all model results, aborts siblings on first + // error, then determines the final outcome by checking context state. + // SIGTERM is NOT checked: output is already flushed to disk, so the caller + // should proceed to finalizeJob rather than re-enqueueing a complete job. + errCh := make(chan error, len(modelMap.SafeToModel)) + resultCh := make(chan error, 1) + go func() { + var firstErr error + for range modelMap.SafeToModel { + if err := <-errCh; err != nil && firstErr == nil { + firstErr = err + abortFn() + } + } + // All model goroutines have finished — safe to signal executeJob to flush. + if firstErr != nil { + resultCh <- firstErr + return + } + switch { + case errors.Is(sloCtx.Err(), context.DeadlineExceeded): + resultCh <- errExpired + case userCancelCtx.Err() != nil: + resultCh <- errCancelled + default: + resultCh <- nil + } + }() + + // Start processing each model. for safeModelID, modelID := range modelMap.SafeToModel { - // Ordering guarantee: processModel returns → requestAbortFn → errCh send. - // This ensures the first real error reaches errCh before any context.Canceled - // from other models whose contexts were cancelled by requestAbortFn. - go func(safeModelID, modelID string) { + go func() { var err error if p.asyncInference != nil { err = p.processModelAsync( @@ -288,8 +290,7 @@ func (p *Processor) executeJob(ctx, sloCtx, userCancelCtx, requestAbortCtx conte userCancelCtx, inputFile, plansDir, safeModelID, modelID, - writers, - progress, + collector, passThroughHeaders, tenantID, ) @@ -301,92 +302,33 @@ func (p *Processor) executeJob(ctx, sloCtx, userCancelCtx, requestAbortCtx conte userCancelCtx, inputFile, plansDir, safeModelID, modelID, - writers, - progress, + collector, passThroughHeaders, tenantID, ) } - // Abort all sibling models when any model hits a fatal I/O error - // (e.g. output file write failure). modelErr is only set for local - // I/O failures — not inference errors, which are recorded normally - // in the error file. Since all models share the same output writers, - // a write failure in one model means the shared file is unusable - // and continuing other models would produce corrupt output. - // Guard against nil requestAbortFn for direct-call test paths. - if err != nil { - if fn := params.requestAbortFn; fn != nil { - fn() - } - } errCh <- err - }(safeModelID, modelID) + }() } - var firstErr error - for range modelMap.SafeToModel { - if err := <-errCh; err != nil && firstErr == nil { - firstErr = err - } - } + // Wait on the result, flush, collect counts, log and exit. + resultErr := <-resultCh - // Push final progress to Redis so the last throttled update doesn't - // leave stale counts visible to polling clients. + collector.flush() progress.flush(ctx) - - if firstErr != nil { - // Flush partial results to disk before routing the error. - // Required for all non-nil firstErr paths: errExpired, errCancelled, system errors - // (callers upload from disk), and SIGTERM (startup recovery reads from disk on restart). - if err := writers.output.Flush(); err != nil { - logger.Error(err, "Failed to flush output file on error path (partial results may be truncated)") - } - if err := writers.errors.Flush(); err != nil { - logger.Error(err, "Failed to flush error file on error path (partial results may be truncated)") - } - // processModel already drained undispatched entries and returned a sentinel (errExpired or - // errCancelled) or the underlying system error. All terminal handlers now use detached - // contexts, so we preserve processModel's decision even when SIGTERM is concurrent. - counts := progress.counts() - switch { - case errors.Is(firstErr, errExpired): - logger.V(logging.INFO).Info("Execution SLO expired, returning partial counts", - "total", counts.Total, "completed", counts.Completed, "failed", counts.Failed) - case errors.Is(firstErr, errCancelled): - logger.V(logging.INFO).Info("Execution cancelled, returning partial counts", - "total", counts.Total, "completed", counts.Completed, "failed", counts.Failed) - default: - logger.V(logging.INFO).Info("Execution system error, returning partial counts", - "total", counts.Total, "completed", counts.Completed, "failed", counts.Failed) - } - return counts, firstErr - } - - if err := writers.output.Flush(); err != nil { - return nil, fmt.Errorf("failed to flush output file: %w", err) - } - if err := writers.errors.Flush(); err != nil { - return nil, fmt.Errorf("failed to flush error file: %w", err) - } - counts := progress.counts() - logger.V(logging.INFO).Info("Execution completed", - "total", counts.Total, "completed", counts.Completed, "failed", counts.Failed) - - // A terminal signal may have arrived after all requests completed normally. - // SLO expiry and user cancel affect the job's terminal status even when all - // requests finished — e.g. "completed but past SLO" vs "completed". - // SIGTERM is NOT checked here: all output is already flushed to disk and counts - // are final, so the caller should proceed to finalizeJob (which uses a detached - // context) rather than re-enqueueing a fully-complete job. - switch { - case errors.Is(sloCtx.Err(), context.DeadlineExceeded): - return counts, errExpired - case userCancelCtx.Err() != nil: - return counts, errCancelled + + var msg string + logVals := []any{"total", counts.Total, "completed", counts.Completed, "failed", counts.Failed} + if resultErr != nil { + msg = "Execution finished" + logVals = append([]any{"error", resultErr}, logVals...) + } else { + msg = "Execution completed" } - return counts, nil + logger.V(logging.INFO).Info(msg, logVals...) + return counts, resultErr } // processModel processes all plan entries for a single model concurrently. @@ -409,8 +351,7 @@ func (p *Processor) processModel( userCancelCtx context.Context, inputFile *os.File, plansDir, safeModelID, modelID string, - writers *outputWriters, - progress *executionProgress, + collector *resultCollector, passThroughHeaders map[string]string, tenantID string, ) error { @@ -434,7 +375,7 @@ func (p *Processor) processModel( epLimit := p.endpointLimits[client] if epLimit == nil { logger.V(logging.INFO).Info("No endpoint limit for model (client not in resolver), draining as model_not_found") - p.drainUnprocessedRequests(requestAbortCtx, inputFile, entries, writers, progress, + p.drainUnprocessedRequests(requestAbortCtx, inputFile, entries, collector, batch_types.BatchErrorCode(inference.ErrCodeModelNotFound)) return nil } @@ -523,20 +464,33 @@ dispatch: return } + // If user-initiated cancel arrived while this request was in-flight, + // overwrite the result as batch_cancelled and send to the collector. + // SLO expiry does not overwrite in-flight results — only user cancel does. + if sloCtx.Err() == nil && userCancelCtx.Err() != nil { + collector.collect(&ResultItem{ + RequestID: result.ID, + CustomID: result.CustomID, + Error: &OutputError{ + Code: string(batch_types.ErrCodeBatchCancelled), + Message: "This request was cancelled while in progress.", + }, + }) + return + } + if result.Error != nil && mainCtx.Err() != nil { shutdownCancelled.Add(1) } - if err := writeResult(result, sloCtx, userCancelCtx, requestAbortCtx, writers, progress); err != nil { - errOnce.Do(func() { modelErr = err }) - } + collector.collect(outputLineToResultItem(result)) }(entry) } wg.Wait() return p.drainAndFinalize(requestAbortCtx, mainCtx, sloCtx, userCancelCtx, - inputFile, entries[dispatchedCount:], writers, progress, modelErr, logger, len(entries), + inputFile, entries[dispatchedCount:], collector, modelErr, logger, len(entries), shutdownCancelled.Load()) } @@ -550,8 +504,7 @@ func (p *Processor) processModelAsync( userCancelCtx context.Context, inputFile *os.File, plansDir, safeModelID, modelID string, - writers *outputWriters, - progress *executionProgress, + collector *resultCollector, passThroughHeaders map[string]string, tenantID string, ) error { @@ -570,7 +523,7 @@ func (p *Processor) processModelAsync( if asyncClient == nil { logger.V(logging.INFO).Info("No async client for model, draining as model_not_found") p.drainUnprocessedRequests( - requestAbortCtx, inputFile, entries, writers, progress, + requestAbortCtx, inputFile, entries, collector, inference.ErrCodeModelNotFound) return nil } @@ -600,15 +553,7 @@ func (p *Processor) processModelAsync( return readErr } if parseErr != nil { - lineBytes, err := json.Marshal(parseErr) - if err != nil { - return fmt.Errorf("marshal parse error line: %w", err) - } - lineBytes = append(lineBytes, '\n') - if err := writers.write(lineBytes, true); err != nil { - return fmt.Errorf("write parse error line: %w", err) - } - progress.record(requestAbortCtx, false) + collector.collect(outputLineToResultItem(parseErr)) submitCount++ continue } @@ -628,17 +573,11 @@ func (p *Processor) processModelAsync( } if submitErr := asyncClient.Submit(requestAbortCtx, inferReq); submitErr != nil { - out := newErrorOutputLine(batchReqID, req.CustomID, - string(submitErr.Category), submitErr.Message) - lineBytes, err := json.Marshal(out) - if err != nil { - return fmt.Errorf("marshal submit error line: %w", err) - } - lineBytes = append(lineBytes, '\n') - if err := writers.write(lineBytes, true); err != nil { - return fmt.Errorf("write submit error line: %w", err) - } - progress.record(requestAbortCtx, false) + collector.collect(&ResultItem{ + RequestID: batchReqID, + CustomID: req.CustomID, + Error: &OutputError{Code: string(submitErr.Category), Message: submitErr.Message}, + }) submitCount++ continue } @@ -672,9 +611,17 @@ func (p *Processor) processModelAsync( } out := buildOutputLine(pr.batchReqID, pr.customID, modelID, resp.RequestID, resp, nil, logger) - if err := writeResult(out, sloCtx, userCancelCtx, requestAbortCtx, writers, progress); err != nil { - modelErr = err - break + if sloCtx.Err() == nil && userCancelCtx.Err() != nil { + collector.collect(&ResultItem{ + RequestID: out.ID, + CustomID: out.CustomID, + Error: &OutputError{ + Code: string(batch_types.ErrCodeBatchCancelled), + Message: "This request was cancelled while in progress.", + }, + }) + } else { + collector.collect(outputLineToResultItem(out)) } delete(pending, resp.RequestID) } @@ -682,21 +629,15 @@ func (p *Processor) processModelAsync( // Drain submitted-but-uncollected requests as errors so that // output_lines + error_lines == total_requests. for _, pr := range pending { - out := newErrorOutputLine(pr.batchReqID, pr.customID, - string(batch_types.ErrCodeBatchExpired), "result not collected before deadline") - lineBytes, err := json.Marshal(out) - if err != nil { - return fmt.Errorf("marshal uncollected error line: %w", err) - } - lineBytes = append(lineBytes, '\n') - if err := writers.write(lineBytes, true); err != nil { - return fmt.Errorf("write uncollected error line: %w", err) - } - progress.record(requestAbortCtx, false) + collector.collect(&ResultItem{ + RequestID: pr.batchReqID, + CustomID: pr.customID, + Error: &OutputError{Code: string(batch_types.ErrCodeBatchExpired), Message: "result not collected before deadline"}, + }) } return p.drainAndFinalize(requestAbortCtx, mainCtx, sloCtx, userCancelCtx, - inputFile, entries[submitCount:], writers, progress, modelErr, logger, len(entries), 0) + inputFile, entries[submitCount:], collector, modelErr, logger, len(entries), 0) } // drainAndFinalize drains undispatched entries based on termination reason and @@ -708,8 +649,7 @@ func (p *Processor) drainAndFinalize( userCancelCtx context.Context, inputFile *os.File, undispatched []planEntry, - writers *outputWriters, - progress *executionProgress, + collector *resultCollector, modelErr error, logger logr.Logger, totalEntries int, @@ -721,7 +661,7 @@ func (p *Processor) drainAndFinalize( // SLO deadline fired during dispatch — record remaining requests as expired. if len(undispatched) > 0 { logger.V(logging.INFO).Info("SLO expired: draining undispatched entries", "count", len(undispatched)) - p.drainUnprocessedRequests(requestAbortCtx, inputFile, undispatched, writers, progress, + p.drainUnprocessedRequests(requestAbortCtx, inputFile, undispatched, collector, batch_types.ErrCodeBatchExpired) } returnErr = errExpired @@ -730,7 +670,7 @@ func (p *Processor) drainAndFinalize( // User-initiated cancel — record remaining requests as cancelled. if len(undispatched) > 0 { logger.V(logging.INFO).Info("Cancelled: draining undispatched entries", "count", len(undispatched)) - p.drainUnprocessedRequests(requestAbortCtx, inputFile, undispatched, writers, progress, + p.drainUnprocessedRequests(requestAbortCtx, inputFile, undispatched, collector, batch_types.ErrCodeBatchCancelled) } returnErr = errCancelled @@ -739,7 +679,7 @@ func (p *Processor) drainAndFinalize( // System error in a model goroutine — record remaining requests as failed. if len(undispatched) > 0 { logger.V(logging.INFO).Info("Fatal error: draining undispatched entries", "count", len(undispatched)) - p.drainUnprocessedRequests(requestAbortCtx, inputFile, undispatched, writers, progress, + p.drainUnprocessedRequests(requestAbortCtx, inputFile, undispatched, collector, batch_types.ErrCodeBatchFailed) } returnErr = modelErr @@ -758,7 +698,7 @@ func (p *Processor) drainAndFinalize( // Drain undispatched entries as batch_failed so that // completed + failed == total holds for the job. logger.V(logging.INFO).Info("Sibling abort: draining undispatched entries", "count", len(undispatched)) - p.drainUnprocessedRequests(requestAbortCtx, inputFile, undispatched, writers, progress, + p.drainUnprocessedRequests(requestAbortCtx, inputFile, undispatched, collector, batch_types.ErrCodeBatchFailed) } } @@ -768,20 +708,18 @@ func (p *Processor) drainAndFinalize( return returnErr } -// drainUnprocessedRequests records undispatched requests in the error file when a job terminates -// mid-execution (SLO expiry, cancellation, or systemic failure). For each plan entry, it reads -// the original request from input.jsonl to extract the custom_id, then writes an error line with -// the given error code and its canonical message. +// drainUnprocessedRequests sends error entries for undispatched plan entries to +// the collector. Called from processModel when dispatch is interrupted (SLO +// expiry, cancellation, or systemic failure). For each entry it reads the +// original request from input.jsonl to extract the custom_id. func (p *Processor) drainUnprocessedRequests( ctx context.Context, inputFile *os.File, entries []planEntry, - writers *outputWriters, - progress *executionProgress, + collector *resultCollector, errCode batch_types.BatchErrorCode, ) { errMessage := errCode.Message() - logger := logr.FromContextOrDiscard(ctx) // Allocate a single read buffer sized to the largest entry to avoid per-entry allocations. var maxLen uint32 @@ -801,25 +739,27 @@ func (p *Processor) drainUnprocessedRequests( } } - requestID := uuid.NewString() - - line := newErrorOutputLine(newBatchRequestID(requestID), customID, string(errCode), errMessage) - - lineBytes, err := json.Marshal(line) - if err != nil { - logger.Error(err, "Failed to marshal drain entry", "errCode", errCode, "offset", entry.Offset) - continue - } - lineBytes = append(lineBytes, '\n') - - if writeErr := writers.write(lineBytes, true); writeErr != nil { - logger.Error(writeErr, "Failed to write drain entry", "errCode", errCode, "offset", entry.Offset) - } + collector.collect(&ResultItem{ + RequestID: newBatchRequestID(uuid.NewString()), + CustomID: customID, + Error: &OutputError{Code: string(errCode), Message: errMessage}, + }) + } +} - // Context may be cancelled here (e.g. SLO deadline fired), so the Redis progress - // update inside record() may fail silently. The atomic counter still increments - // correctly and the final counts are committed by the terminal status update. - progress.record(ctx, false) +// outputLineToResultItem converts an outputLine (returned by executeOneRequest) to a +// ResultItem for the collector. +func outputLineToResultItem(ol *outputLine) *ResultItem { + var outErr *OutputError + if ol.Error != nil { + outErr = &OutputError{Code: ol.Error.Code, Message: ol.Error.Message} + } + return &ResultItem{ + RequestID: ol.ID, + CustomID: ol.CustomID, + Response: ol.Response, + Error: outErr, + HadCapacityRetry: ol.hadCapacityRetry, } } @@ -983,35 +923,6 @@ func newErrorOutputLine(batchReqID, customID, code, message string) *outputLine } } -// writeResult applies user-cancel overwrite if needed, records progress, marshals -// the output line, and writes it to the appropriate file. Returns an error only -// for marshal/write failures. -func writeResult( - out *outputLine, - sloCtx, userCancelCtx, progressCtx context.Context, - writers *outputWriters, - progress *executionProgress, -) error { - if sloCtx.Err() == nil && userCancelCtx.Err() != nil { - out.Response = nil - out.Error = &outputError{ - Code: string(batch_types.ErrCodeBatchCancelled), - Message: "This request was cancelled while in progress.", - } - } - - progress.record(progressCtx, out.isSuccess()) - lineBytes, err := json.Marshal(out) - if err != nil { - return fmt.Errorf("marshal output line for %s: %w", out.ID, err) - } - lineBytes = append(lineBytes, '\n') - if writeErr := writers.write(lineBytes, out.Error != nil); writeErr != nil { - return fmt.Errorf("write output line for %s: %w", out.ID, writeErr) - } - return nil -} - // buildOutputLine converts an inference response and/or error into an outputLine. // Used by both executeOneRequest (sync path) and processModelAsync (async path). func buildOutputLine( diff --git a/internal/processor/worker/executor_test.go b/internal/processor/worker/executor_test.go index 83d736b0d..f5c9a6bc1 100644 --- a/internal/processor/worker/executor_test.go +++ b/internal/processor/worker/executor_test.go @@ -802,27 +802,22 @@ func TestProcessModel_Success(t *testing.T) { plansDir, _ := env.p.jobPlansDir(jobInfo.JobID, jobInfo.TenantID) - var buf bytes.Buffer - writer := bufio.NewWriter(&buf) - + var outBuf, errBuf bytes.Buffer progress := &executionProgress{ total: int64(len(requests)), updater: env.updater, jobID: jobInfo.JobID, } - - var errBuf bytes.Buffer - writers := &outputWriters{output: writer, errors: bufio.NewWriter(&errBuf)} + collector := newResultCollector(bufio.NewWriter(&outBuf), bufio.NewWriter(&errBuf), progress, logr.Discard(), func() {}) + collector.start(testLoggerCtx(t)) ctx := testLoggerCtx(t) - err := env.p.processModel(ctx, ctx, ctx, context.Background(), inputFile, plansDir, "m1", "m1", writers, progress, nil, "") + err := env.p.processModel(ctx, ctx, ctx, context.Background(), inputFile, plansDir, "m1", "m1", collector, nil, "") if err != nil { t.Fatalf("processModel error: %v", err) } - if err := writer.Flush(); err != nil { - t.Fatalf("flush: %v", err) - } + collector.flush() if int(callCount.Load()) != len(requests) { t.Fatalf("inference calls = %d, want %d", callCount.Load(), len(requests)) @@ -833,7 +828,7 @@ func TestProcessModel_Success(t *testing.T) { t.Fatalf("completed = %d, want %d", counts.Completed, len(requests)) } - lines := bytes.Split(bytes.TrimSpace(buf.Bytes()), []byte{'\n'}) + lines := bytes.Split(bytes.TrimSpace(outBuf.Bytes()), []byte{'\n'}) if len(lines) != len(requests) { t.Fatalf("output lines = %d, want %d", len(lines), len(requests)) } @@ -856,18 +851,14 @@ func TestProcessModel_CancelStopsDispatch(t *testing.T) { plansDir, _ := env.p.jobPlansDir(jobInfo.JobID, jobInfo.TenantID) - var buf bytes.Buffer - writer := bufio.NewWriter(&buf) - + var outBuf, errBuf bytes.Buffer progress := &executionProgress{ total: 1, updater: env.updater, jobID: jobInfo.JobID, } - - var errBuf bytes.Buffer - errWriter := bufio.NewWriter(&errBuf) - writers := &outputWriters{output: writer, errors: errWriter} + collector := newResultCollector(bufio.NewWriter(&outBuf), bufio.NewWriter(&errBuf), progress, logr.Discard(), func() {}) + collector.start(testLoggerCtx(t)) // Cancel ctx to simulate requestAbortCtx being cancelled (by watchCancel calling requestAbortFn). // Separately pass ctx as userCancelCtx so drain chooses errCancelled, not errShutdown. @@ -877,15 +868,13 @@ func TestProcessModel_CancelStopsDispatch(t *testing.T) { ctx, cancel := context.WithCancel(baseCtx) cancel() - err := env.p.processModel(ctx, baseCtx, context.Background(), ctx, inputFile, plansDir, "m1", "m1", writers, progress, nil, "") + err := env.p.processModel(ctx, baseCtx, context.Background(), ctx, inputFile, plansDir, "m1", "m1", collector, nil, "") if !errors.Is(err, errCancelled) { t.Fatalf("expected errCancelled, got: %v", err) } // Verify that undispatched entry was drained as batch_cancelled. - if flushErr := errWriter.Flush(); flushErr != nil { - t.Fatalf("flush error writer: %v", flushErr) - } + collector.flush() errLines := bytes.Split(bytes.TrimSpace(errBuf.Bytes()), []byte{'\n'}) if len(errLines) != 1 { t.Fatalf("expected 1 drain entry in error output, got %d", len(errLines)) @@ -931,28 +920,21 @@ func TestProcessModel_CancelWritesInFlightToErrorFile(t *testing.T) { plansDir, _ := env.p.jobPlansDir(jobInfo.JobID, jobInfo.TenantID) var outBuf, errBuf bytes.Buffer - outWriter := bufio.NewWriter(&outBuf) - errWriter := bufio.NewWriter(&errBuf) - writers := &outputWriters{output: outWriter, errors: errWriter} - progress := &executionProgress{ total: 1, updater: env.updater, jobID: jobInfo.JobID, } + collector := newResultCollector(bufio.NewWriter(&outBuf), bufio.NewWriter(&errBuf), progress, logr.Discard(), func() {}) + collector.start(testLoggerCtx(t)) ctx := testLoggerCtx(t) - modelErr := env.p.processModel(ctx, ctx, ctx, userCancelCtx, inputFile, plansDir, "m1", "m1", writers, progress, nil, "") + modelErr := env.p.processModel(ctx, ctx, ctx, userCancelCtx, inputFile, plansDir, "m1", "m1", collector, nil, "") if !errors.Is(modelErr, errCancelled) { t.Fatalf("expected errCancelled from processModel, got: %v", modelErr) } - if flushErr := outWriter.Flush(); flushErr != nil { - t.Fatalf("flush output: %v", flushErr) - } - if flushErr := errWriter.Flush(); flushErr != nil { - t.Fatalf("flush error: %v", flushErr) - } + collector.flush() // Output file should be empty — cancelled requests go to error file. if outBuf.Len() > 0 { @@ -1011,23 +993,21 @@ func TestProcessModel_InferenceFatalError(t *testing.T) { plansDir, _ := env.p.jobPlansDir(jobInfo.JobID, jobInfo.TenantID) - var buf bytes.Buffer - writer := bufio.NewWriter(&buf) - + var outBuf, errBuf bytes.Buffer progress := &executionProgress{ total: int64(len(requests)), updater: env.updater, jobID: jobInfo.JobID, } - - var errBuf bytes.Buffer - writers := &outputWriters{output: writer, errors: bufio.NewWriter(&errBuf)} + collector := newResultCollector(bufio.NewWriter(&outBuf), bufio.NewWriter(&errBuf), progress, logr.Discard(), func() {}) + collector.start(testLoggerCtx(t)) ctx := testLoggerCtx(t) - err := env.p.processModel(ctx, ctx, ctx, context.Background(), inputFile, plansDir, "m1", "m1", writers, progress, nil, "") + err := env.p.processModel(ctx, ctx, ctx, context.Background(), inputFile, plansDir, "m1", "m1", collector, nil, "") if err == nil { t.Fatalf("expected error from closed input file") } + collector.flush() } func TestProcessModel_ContextCancelledDuringDispatch(t *testing.T) { @@ -1058,23 +1038,20 @@ func TestProcessModel_ContextCancelledDuringDispatch(t *testing.T) { plansDir, _ := env.p.jobPlansDir(jobInfo.JobID, jobInfo.TenantID) - var buf bytes.Buffer - writer := bufio.NewWriter(&buf) - + var outBuf, errBuf bytes.Buffer progress := &executionProgress{ total: int64(len(requests)), updater: env.updater, jobID: jobInfo.JobID, } + collector := newResultCollector(bufio.NewWriter(&outBuf), bufio.NewWriter(&errBuf), progress, logr.Discard(), func() {}) + collector.start(testLoggerCtx(t)) ctx, cancel := context.WithCancel(testLoggerCtx(t)) - var errBuf bytes.Buffer - writers := &outputWriters{output: writer, errors: bufio.NewWriter(&errBuf)} - done := make(chan error, 1) go func() { - done <- env.p.processModel(ctx, ctx, ctx, context.Background(), inputFile, plansDir, "m1", "m1", writers, progress, nil, "") + done <- env.p.processModel(ctx, ctx, ctx, context.Background(), inputFile, plansDir, "m1", "m1", collector, nil, "") }() <-started @@ -1082,6 +1059,7 @@ func TestProcessModel_ContextCancelledDuringDispatch(t *testing.T) { close(block) err := <-done + collector.flush() if err == nil { t.Fatalf("expected error on context cancellation") } @@ -1123,18 +1101,19 @@ func TestProcessModel_SIGTERMCancelsAllDispatched(t *testing.T) { plansDir, _ := env.p.jobPlansDir(jobInfo.JobID, jobInfo.TenantID) var outBuf, errBuf bytes.Buffer - writers := &outputWriters{output: bufio.NewWriter(&outBuf), errors: bufio.NewWriter(&errBuf)} - progress := &executionProgress{ total: int64(len(requests)), updater: env.updater, jobID: jobInfo.JobID, } + collector := newResultCollector(bufio.NewWriter(&outBuf), bufio.NewWriter(&errBuf), progress, logr.Discard(), func() {}) + collector.start(testLoggerCtx(t)) - err := env.p.processModel(mainCtx, mainCtx, mainCtx, context.Background(), inputFile, plansDir, "m1", "m1", writers, progress, nil, "") + err := env.p.processModel(mainCtx, mainCtx, mainCtx, context.Background(), inputFile, plansDir, "m1", "m1", collector, nil, "") if !errors.Is(err, errShutdown) { t.Fatalf("expected errShutdown when SIGTERM cancels all dispatched requests, got: %v", err) } + collector.flush() } // TestProcessModel_SiblingAbort_ReturnsNil verifies that when requestAbortCtx is cancelled @@ -1160,14 +1139,14 @@ func TestProcessModel_SiblingAbort_ReturnsNil(t *testing.T) { plansDir, _ := env.p.jobPlansDir(jobInfo.JobID, jobInfo.TenantID) - var buf bytes.Buffer - writers := &outputWriters{output: bufio.NewWriter(&buf), errors: bufio.NewWriter(&buf)} - + var outBuf, errBuf bytes.Buffer progress := &executionProgress{ total: 1, updater: env.updater, jobID: jobInfo.JobID, } + collector := newResultCollector(bufio.NewWriter(&outBuf), bufio.NewWriter(&errBuf), progress, logr.Discard(), func() {}) + collector.start(testLoggerCtx(t)) // mainCtx is not cancelled — only requestAbortCtx is, simulating a sibling model calling // requestAbortFn() on error. SLO and user-cancel signals are both absent. @@ -1175,11 +1154,12 @@ func TestProcessModel_SiblingAbort_ReturnsNil(t *testing.T) { requestAbortCtx, requestAbortFn := context.WithCancel(mainCtx) requestAbortFn() // simulate sibling model calling requestAbortFn - err := env.p.processModel(requestAbortCtx, mainCtx, mainCtx, context.Background(), inputFile, plansDir, "m1", "m1", writers, progress, nil, "") + err := env.p.processModel(requestAbortCtx, mainCtx, mainCtx, context.Background(), inputFile, plansDir, "m1", "m1", collector, nil, "") // requestAbortCtx cancelled, but no SLO / user-cancel / SIGTERM → nil, not errShutdown if err != nil { t.Fatalf("expected nil when only requestAbortCtx is cancelled (sibling abort), got: %v", err) } + collector.flush() } // ===================================================================== @@ -3389,21 +3369,20 @@ func TestProcessModel_AIMDSignaling(t *testing.T) { plansDir, _ := env.p.jobPlansDir(jobInfo.JobID, jobInfo.TenantID) var outBuf, errBuf bytes.Buffer - writers := &outputWriters{ - output: bufio.NewWriter(&outBuf), - errors: bufio.NewWriter(&errBuf), - } progress := &executionProgress{ total: int64(len(requests)), updater: env.updater, jobID: jobInfo.JobID, } + collector := newResultCollector(bufio.NewWriter(&outBuf), bufio.NewWriter(&errBuf), progress, logr.Discard(), func() {}) + collector.start(testLoggerCtx(t)) ctx := testLoggerCtx(t) - err = env.p.processModel(ctx, ctx, ctx, context.Background(), inputFile, plansDir, "m1", "m1", writers, progress, nil, jobInfo.TenantID) + err = env.p.processModel(ctx, ctx, ctx, context.Background(), inputFile, plansDir, "m1", "m1", collector, nil, jobInfo.TenantID) if err != nil { t.Fatalf("processModel error: %v", err) } + collector.flush() return env.p } @@ -3548,21 +3527,20 @@ func TestProcessModel_AIMDSignaling(t *testing.T) { plansDir, _ := env.p.jobPlansDir(jobInfo.JobID, jobInfo.TenantID) var outBuf, errBuf bytes.Buffer - writers := &outputWriters{ - output: bufio.NewWriter(&outBuf), - errors: bufio.NewWriter(&errBuf), - } progress := &executionProgress{ total: int64(len(requests)), updater: env.updater, jobID: jobInfo.JobID, } + collector := newResultCollector(bufio.NewWriter(&outBuf), bufio.NewWriter(&errBuf), progress, logr.Discard(), func() {}) + collector.start(testLoggerCtx(t)) ctx := testLoggerCtx(t) - err = env.p.processModel(ctx, ctx, ctx, context.Background(), inputFile, plansDir, "m1", "m1", writers, progress, nil, jobInfo.TenantID) + err = env.p.processModel(ctx, ctx, ctx, context.Background(), inputFile, plansDir, "m1", "m1", collector, nil, jobInfo.TenantID) if err != nil { t.Fatalf("processModel error: %v", err) } + collector.flush() // If non-HTTP errors were incorrectly counted as RecordSuccess, // the window (size=reducedLimit) would fill and push the limit @@ -3670,18 +3648,18 @@ func TestProcessModel_AIMDEndpointIsolation(t *testing.T) { updater := NewStatusUpdater(dbClient, statusClient, 86400) ctx := testLoggerCtx(t) - // Process m1 (429s) — should decrease m1's endpoint AIMD. var outBuf, errBuf bytes.Buffer - writers := &outputWriters{ - output: bufio.NewWriter(&outBuf), - errors: bufio.NewWriter(&errBuf), - } progress := &executionProgress{total: 4, updater: updater, jobID: jobID} + collector := newResultCollector(bufio.NewWriter(&outBuf), bufio.NewWriter(&errBuf), progress, logr.Discard(), func() {}) + collector.start(testLoggerCtx(t)) - _ = p.processModel(ctx, ctx, ctx, context.Background(), inputFile, plansDir, "m1", "m1", writers, progress, nil, tenantID) + // Process m1 (429s) — should decrease m1's endpoint AIMD. + _ = p.processModel(ctx, ctx, ctx, context.Background(), inputFile, plansDir, "m1", "m1", collector, nil, tenantID) // Process m2 (200s) — should NOT affect m2's endpoint AIMD. - _ = p.processModel(ctx, ctx, ctx, context.Background(), inputFile, plansDir, "m2", "m2", writers, progress, nil, tenantID) + _ = p.processModel(ctx, ctx, ctx, context.Background(), inputFile, plansDir, "m2", "m2", collector, nil, tenantID) + + collector.flush() limitA := p.endpointLimits[clientA].aimd.Limit() limitB := p.endpointLimits[clientB].aimd.Limit() @@ -3732,28 +3710,21 @@ func TestProcessModel_EndpointLimitNil_DrainsAsModelNotFound(t *testing.T) { plansDir, _ := env.p.jobPlansDir(jobInfo.JobID, jobInfo.TenantID) var outBuf, errBuf bytes.Buffer - writers := &outputWriters{ - output: bufio.NewWriter(&outBuf), - errors: bufio.NewWriter(&errBuf), - } progress := &executionProgress{ total: int64(len(requests)), updater: env.updater, jobID: jobInfo.JobID, } + collector := newResultCollector(bufio.NewWriter(&outBuf), bufio.NewWriter(&errBuf), progress, logr.Discard(), func() {}) + collector.start(testLoggerCtx(t)) ctx := testLoggerCtx(t) - err = env.p.processModel(ctx, ctx, ctx, context.Background(), inputFile, plansDir, "m1", "m1", writers, progress, nil, jobInfo.TenantID) + err = env.p.processModel(ctx, ctx, ctx, context.Background(), inputFile, plansDir, "m1", "m1", collector, nil, jobInfo.TenantID) if err != nil { t.Fatalf("processModel error: %v", err) } - if err := writers.errors.Flush(); err != nil { - t.Fatalf("flush errors: %v", err) - } - if err := writers.output.Flush(); err != nil { - t.Fatalf("flush output: %v", err) - } + collector.flush() // All requests should appear in the error file as model_not_found. errLines := strings.Split(strings.TrimSpace(errBuf.String()), "\n") @@ -3912,17 +3883,18 @@ func TestProcessModelAsync(t *testing.T) { plansDir, _ := env.p.jobPlansDir(jobInfo.JobID, jobInfo.TenantID) var outBuf, errBuf bytes.Buffer - writers := &outputWriters{output: bufio.NewWriter(&outBuf), errors: bufio.NewWriter(&errBuf)} progress := &executionProgress{total: int64(len(requests)), updater: env.updater, jobID: jobInfo.JobID} ctx := testLoggerCtx(t) - err := env.p.processModelAsync(ctx, ctx, ctx, context.Background(), inputFile, plansDir, "m1", "m1", writers, progress, nil, "") + collector := newResultCollector(bufio.NewWriter(&outBuf), bufio.NewWriter(&errBuf), progress, logr.FromContextOrDiscard(ctx), func() {}) + collector.start(ctx) + + err := env.p.processModelAsync(ctx, ctx, ctx, context.Background(), inputFile, plansDir, "m1", "m1", collector, nil, "") if err != nil { t.Fatalf("processModelAsync error: %v", err) } - _ = writers.output.Flush() - _ = writers.errors.Flush() + collector.flush() if len(submitted) != len(requests) { t.Fatalf("submitted = %d, want %d", len(submitted), len(requests)) @@ -3983,16 +3955,18 @@ func TestProcessModelAsync(t *testing.T) { plansDir, _ := env.p.jobPlansDir(jobInfo.JobID, jobInfo.TenantID) var outBuf, errBuf bytes.Buffer - writers := &outputWriters{output: bufio.NewWriter(&outBuf), errors: bufio.NewWriter(&errBuf)} progress := &executionProgress{total: int64(len(requests)), updater: env.updater, jobID: jobInfo.JobID} ctx := testLoggerCtx(t) abortCtx, abortFn := context.WithCancel(ctx) + collector := newResultCollector(bufio.NewWriter(&outBuf), bufio.NewWriter(&errBuf), progress, logr.FromContextOrDiscard(ctx), func() {}) + collector.start(ctx) + // Run processModelAsync in a goroutine so we can intercept submitted IDs. done := make(chan error, 1) go func() { - done <- env.p.processModelAsync(abortCtx, ctx, ctx, context.Background(), inputFile, plansDir, "m1", "m1", writers, progress, nil, "") + done <- env.p.processModelAsync(abortCtx, ctx, ctx, context.Background(), inputFile, plansDir, "m1", "m1", collector, nil, "") }() // Wait for all 3 submits, deliver 1 result using the real ID, then cancel. @@ -4010,8 +3984,7 @@ func TestProcessModelAsync(t *testing.T) { <-done - _ = writers.output.Flush() - _ = writers.errors.Flush() + collector.flush() counts := progress.counts() // 1 completed + 2 expired = 3 total @@ -4045,17 +4018,18 @@ func TestProcessModelAsync(t *testing.T) { plansDir, _ := env.p.jobPlansDir(jobInfo.JobID, jobInfo.TenantID) var outBuf, errBuf bytes.Buffer - writers := &outputWriters{output: bufio.NewWriter(&outBuf), errors: bufio.NewWriter(&errBuf)} progress := &executionProgress{total: 1, updater: env.updater, jobID: jobInfo.JobID} ctx := testLoggerCtx(t) - err := env.p.processModelAsync(ctx, ctx, ctx, context.Background(), inputFile, plansDir, "m1", "m1", writers, progress, nil, "") + collector := newResultCollector(bufio.NewWriter(&outBuf), bufio.NewWriter(&errBuf), progress, logr.FromContextOrDiscard(ctx), func() {}) + collector.start(ctx) + + err := env.p.processModelAsync(ctx, ctx, ctx, context.Background(), inputFile, plansDir, "m1", "m1", collector, nil, "") if err != nil { t.Fatalf("processModelAsync error: %v", err) } - _ = writers.output.Flush() - _ = writers.errors.Flush() + collector.flush() counts := progress.counts() if counts.Failed != 1 { From 762e139a164576508e7cf0e1119b0d5153210405 Mon Sep 17 00:00:00 2001 From: Edoardo Vacchi Date: Thu, 9 Jul 2026 17:52:08 +0200 Subject: [PATCH 2/2] address comments Signed-off-by: Edoardo Vacchi --- internal/processor/worker/collector.go | 33 ++++++++----- internal/processor/worker/executor.go | 4 +- internal/processor/worker/executor_test.go | 56 ++++++++++++++++------ 3 files changed, 65 insertions(+), 28 deletions(-) diff --git a/internal/processor/worker/collector.go b/internal/processor/worker/collector.go index 68a7e9d47..4c213836f 100644 --- a/internal/processor/worker/collector.go +++ b/internal/processor/worker/collector.go @@ -83,8 +83,9 @@ type resultCollector struct { abortFn context.CancelFunc abortOnce sync.Once - ch chan *ResultItem - done chan struct{} + ch chan *ResultItem + done chan error + writeErr error } func newResultCollector(outputWriter, errorWriter *bufio.Writer, progress *executionProgress, logger logr.Logger, abortFn context.CancelFunc) *resultCollector { @@ -98,7 +99,7 @@ func newResultCollector(outputWriter, errorWriter *bufio.Writer, progress *execu logger: logger, abortFn: abortFn, ch: make(chan *ResultItem, 1024), - done: make(chan struct{}), + done: make(chan error, 1), } } @@ -106,8 +107,7 @@ func newResultCollector(outputWriter, errorWriter *bufio.Writer, progress *execu // updates to the status store. func (c *resultCollector) start(ctx context.Context) { go func() { - c.run(ctx) - close(c.done) + c.done <- c.run(ctx) }() } @@ -116,25 +116,29 @@ func (c *resultCollector) collect(result *ResultItem) { c.ch <- result } -// flush closes the result channel and waits for the collector goroutine to -// finish writing all buffered results and flushing the underlying writers. -func (c *resultCollector) flush() { +// flush closes the result channel, waits for the collector goroutine to +// finish writing all buffered results, and returns the first write error +// encountered (or nil). +func (c *resultCollector) flush() error { close(c.ch) - <-c.done + return <-c.done } -func (c *resultCollector) abort() { +func (c *resultCollector) abort(err error) { + if c.writeErr == nil { + c.writeErr = err + } c.abortOnce.Do(c.abortFn) } -func (c *resultCollector) run(ctx context.Context) { +func (c *resultCollector) run(ctx context.Context) error { for result := range c.ch { line := resultToOutputLine(result) lineBytes, err := json.Marshal(line) if err != nil { c.logger.Error(err, "Failed to marshal output line", "customId", result.CustomID) - c.abort() + c.abort(err) continue } lineBytes = append(lineBytes, '\n') @@ -146,7 +150,7 @@ func (c *resultCollector) run(ctx context.Context) { } if _, err := writer.Write(lineBytes); err != nil { c.logger.Error(err, "Failed to write output line", "customId", result.CustomID) - c.abort() + c.abort(err) continue } @@ -155,8 +159,11 @@ func (c *resultCollector) run(ctx context.Context) { if err := c.outputWriter.Flush(); err != nil { c.logger.Error(err, "Failed to flush output file (partial results may be truncated)") + c.abort(err) } if err := c.errorWriter.Flush(); err != nil { c.logger.Error(err, "Failed to flush error file (partial results may be truncated)") + c.abort(err) } + return c.writeErr } diff --git a/internal/processor/worker/executor.go b/internal/processor/worker/executor.go index be6fba7f2..fa762903d 100644 --- a/internal/processor/worker/executor.go +++ b/internal/processor/worker/executor.go @@ -314,7 +314,9 @@ func (p *Processor) executeJob(ctx, sloCtx, userCancelCtx, requestAbortCtx conte // Wait on the result, flush, collect counts, log and exit. resultErr := <-resultCh - collector.flush() + if flushErr := collector.flush(); resultErr == nil { + resultErr = flushErr + } progress.flush(ctx) counts := progress.counts() diff --git a/internal/processor/worker/executor_test.go b/internal/processor/worker/executor_test.go index f5c9a6bc1..675356188 100644 --- a/internal/processor/worker/executor_test.go +++ b/internal/processor/worker/executor_test.go @@ -817,7 +817,9 @@ func TestProcessModel_Success(t *testing.T) { t.Fatalf("processModel error: %v", err) } - collector.flush() + if err := collector.flush(); err != nil { + t.Fatalf("collector.flush: %v", err) + } if int(callCount.Load()) != len(requests) { t.Fatalf("inference calls = %d, want %d", callCount.Load(), len(requests)) @@ -874,7 +876,9 @@ func TestProcessModel_CancelStopsDispatch(t *testing.T) { } // Verify that undispatched entry was drained as batch_cancelled. - collector.flush() + if err := collector.flush(); err != nil { + t.Fatalf("collector.flush: %v", err) + } errLines := bytes.Split(bytes.TrimSpace(errBuf.Bytes()), []byte{'\n'}) if len(errLines) != 1 { t.Fatalf("expected 1 drain entry in error output, got %d", len(errLines)) @@ -934,7 +938,9 @@ func TestProcessModel_CancelWritesInFlightToErrorFile(t *testing.T) { t.Fatalf("expected errCancelled from processModel, got: %v", modelErr) } - collector.flush() + if err := collector.flush(); err != nil { + t.Fatalf("collector.flush: %v", err) + } // Output file should be empty — cancelled requests go to error file. if outBuf.Len() > 0 { @@ -1007,7 +1013,9 @@ func TestProcessModel_InferenceFatalError(t *testing.T) { if err == nil { t.Fatalf("expected error from closed input file") } - collector.flush() + if err := collector.flush(); err != nil { + t.Fatalf("collector.flush: %v", err) + } } func TestProcessModel_ContextCancelledDuringDispatch(t *testing.T) { @@ -1059,7 +1067,9 @@ func TestProcessModel_ContextCancelledDuringDispatch(t *testing.T) { close(block) err := <-done - collector.flush() + if err := collector.flush(); err != nil { + t.Fatalf("collector.flush: %v", err) + } if err == nil { t.Fatalf("expected error on context cancellation") } @@ -1113,7 +1123,9 @@ func TestProcessModel_SIGTERMCancelsAllDispatched(t *testing.T) { if !errors.Is(err, errShutdown) { t.Fatalf("expected errShutdown when SIGTERM cancels all dispatched requests, got: %v", err) } - collector.flush() + if err := collector.flush(); err != nil { + t.Fatalf("collector.flush: %v", err) + } } // TestProcessModel_SiblingAbort_ReturnsNil verifies that when requestAbortCtx is cancelled @@ -1159,7 +1171,9 @@ func TestProcessModel_SiblingAbort_ReturnsNil(t *testing.T) { if err != nil { t.Fatalf("expected nil when only requestAbortCtx is cancelled (sibling abort), got: %v", err) } - collector.flush() + if err := collector.flush(); err != nil { + t.Fatalf("collector.flush: %v", err) + } } // ===================================================================== @@ -3382,7 +3396,9 @@ func TestProcessModel_AIMDSignaling(t *testing.T) { if err != nil { t.Fatalf("processModel error: %v", err) } - collector.flush() + if err := collector.flush(); err != nil { + t.Fatalf("collector.flush: %v", err) + } return env.p } @@ -3540,7 +3556,9 @@ func TestProcessModel_AIMDSignaling(t *testing.T) { if err != nil { t.Fatalf("processModel error: %v", err) } - collector.flush() + if err := collector.flush(); err != nil { + t.Fatalf("collector.flush: %v", err) + } // If non-HTTP errors were incorrectly counted as RecordSuccess, // the window (size=reducedLimit) would fill and push the limit @@ -3659,7 +3677,9 @@ func TestProcessModel_AIMDEndpointIsolation(t *testing.T) { // Process m2 (200s) — should NOT affect m2's endpoint AIMD. _ = p.processModel(ctx, ctx, ctx, context.Background(), inputFile, plansDir, "m2", "m2", collector, nil, tenantID) - collector.flush() + if err := collector.flush(); err != nil { + t.Fatalf("collector.flush: %v", err) + } limitA := p.endpointLimits[clientA].aimd.Limit() limitB := p.endpointLimits[clientB].aimd.Limit() @@ -3724,7 +3744,9 @@ func TestProcessModel_EndpointLimitNil_DrainsAsModelNotFound(t *testing.T) { t.Fatalf("processModel error: %v", err) } - collector.flush() + if err := collector.flush(); err != nil { + t.Fatalf("collector.flush: %v", err) + } // All requests should appear in the error file as model_not_found. errLines := strings.Split(strings.TrimSpace(errBuf.String()), "\n") @@ -3894,7 +3916,9 @@ func TestProcessModelAsync(t *testing.T) { t.Fatalf("processModelAsync error: %v", err) } - collector.flush() + if err := collector.flush(); err != nil { + t.Fatalf("collector.flush: %v", err) + } if len(submitted) != len(requests) { t.Fatalf("submitted = %d, want %d", len(submitted), len(requests)) @@ -3984,7 +4008,9 @@ func TestProcessModelAsync(t *testing.T) { <-done - collector.flush() + if err := collector.flush(); err != nil { + t.Fatalf("collector.flush: %v", err) + } counts := progress.counts() // 1 completed + 2 expired = 3 total @@ -4029,7 +4055,9 @@ func TestProcessModelAsync(t *testing.T) { t.Fatalf("processModelAsync error: %v", err) } - collector.flush() + if err := collector.flush(); err != nil { + t.Fatalf("collector.flush: %v", err) + } counts := progress.counts() if counts.Failed != 1 {