Skip to content

Commit 7828eb5

Browse files
lioraronlioraron
authored
feat: orphan reconciler for stuck batch jobs (llm-d#452)
* feat: orphan reconciler, NonTerminal query, GC wiring, heartbeat logging Add an orphan reconciler that periodically cross-references non-terminal jobs (DB), queued jobs (priority queue), and in-flight entries to detect and recover orphaned batch jobs. Triage logic transitions orphans based on their current status and SLO expiry using CAS updates to prevent overwriting concurrent processor decisions. Wire the reconciler into batch-gc alongside the existing GC collector via errgroup. Add NonTerminal query support to BatchQuery for both PostgreSQL (raw SQL condition) and Redis (new Lua script). Add heartbeat start/refresh/stop logging to the processor worker. Closes llm-d#434 Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com> Signed-off-by: Lior Aronovich <lioraronpr@gmail.com> * fix: address code review findings for orphan reconciler - Add NewReconciler input validation (nil clients, non-positive interval) - Separate CAS conflicts from errors in Result (new Conflicts field) - Extract TerminalStatuses() into shared openai package - Convert extractSLO/isSLOExpired to package-level functions - Fix unreachable code and missing error handling in triage paths - Add comprehensive tests: edge cases, validation, terminal status sync - Add Redis NonTerminal query integration test - Add MockInFlightClient.SetLastSeen test helper for stale entry testing Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com> Signed-off-by: Lior Aronovich <lioraronpr@gmail.com> * feat: add dry-run mode to orphan reconciler When dry-run is enabled, the reconciler detects and counts orphans but skips all mutating operations (DB transitions, queue enqueues, in-flight deletes). The reconciler inherits the GC's existing dry_run config flag. Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com> Signed-off-by: Lior Aronovich <lioraronpr@gmail.com> * feat: heartbeat DB status check to detect reconciler actions On each heartbeat tick, the processor now reads the job's DB status. If the reconciler has acted (terminal status or reverted to validating), the heartbeat calls requestAbortFn to stop all in-flight requests. The processor's subsequent CAS write fails with ErrConflict, and it yields cleanly. This cuts wasted processing from up to 60 minutes (full reconciler interval) to ~5 minutes (heartbeat interval). Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com> Signed-off-by: Lior Aronovich <lioraronpr@gmail.com> * refactor: precompute NonTerminal SQL condition, move NonTerminal into dbGet - PostgreSQL: compute the NOT IN clause once at package init via buildNonTerminalCondition() instead of rebuilding on every DBGet call - Redis: move NonTerminal Lua script branch into dbGet alongside the existing purpose/expiry/tenant branches for consistency - Add comment on TerminalStatuses() explaining per-call allocation Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com> Signed-off-by: Lior Aronovich <lioraronpr@gmail.com> * refactor: rename NonTerminal Lua script to align with naming convention Rename redis_get_non_terminal.lua → redis_get_by_non_terminal.lua and update Go variable names (getNonTerminalLua → getByNonTerminalLua, redisScriptGetNonTerminal → redisScriptGetByNonTerminal) to match the existing redisScriptGetBy* pattern. Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com> Signed-off-by: Lior Aronovich <lioraronpr@gmail.com> * refactor: replace goto/label with nested conditionals in NonTerminal Lua script Consistent with the other Redis Lua scripts in this package which use nested if blocks instead of goto. Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com> Signed-off-by: Lior Aronovich <lioraronpr@gmail.com> * refactor: rename IsFinal to IsTerminal, consistent heartbeat log prefix - Rename BatchStatus.IsFinal() → IsTerminal() to align with TerminalStatuses() naming - Use logr.WithValues for jobId in heartbeat logger - Consistent "Heartbeat: ..." prefix on all heartbeat log messages - Pass logger to checkReconcilerActed to avoid redundant WithValues Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com> Signed-off-by: Lior Aronovich <lioraronpr@gmail.com> * refactor: add Reconciler: log prefix, remove unused run() return, flip dry-run guard Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com> Signed-off-by: Lior Aronovich <lioraronpr@gmail.com> * fix: add batchDB nil check in validate(), skip in-flight delete on failed triage - Processor.validate() now checks batchDB for nil, preventing a panic in checkReconcilerActed during heartbeat. - transitionOrphan and reEnqueueOrphan now return bool to signal success. triageOrphan only deletes the in-flight entry when the action succeeded, avoiding removal of a legitimate processor's heartbeat tracking after a CAS conflict. Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com> Signed-off-by: Lior Aronovich <lioraronpr@gmail.com> * refactor: use shared terminalStatuses array in IsTerminal and TerminalStatuses Single source of truth for terminal statuses. Declared as an array so it cannot be appended to or resliced. IsTerminal ranges over it directly (no allocation). TerminalStatuses returns a copy so callers cannot mutate the canonical list. Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com> Signed-off-by: Lior Aronovich <lioraronpr@gmail.com> * refactor: skip queue query when no non-terminal jobs exist When fetchNonTerminalJobs returns zero jobs, skip PQGetIDs and the triage loop. InFlightGetAll and cleanupStaleInflight still run so stale entries for already-terminal jobs are cleaned up. Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com> Signed-off-by: Lior Aronovich <lioraronpr@gmail.com> * refactor: simplify batch-gc main to always use errgroup Removes the duplicated gc.RunLoop path when reconciler is disabled. The errgroup always runs the GC collector; the reconciler is conditionally added when enabled. Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com> Signed-off-by: Lior Aronovich <lioraronpr@gmail.com> --------- Signed-off-by: Lior Aronovich <lioraronpr@gmail.com>
1 parent 657d95e commit 7828eb5

21 files changed

Lines changed: 1585 additions & 25 deletions

‎cmd/batch-gc/main.go‎

Lines changed: 23 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -29,8 +29,11 @@ import (
2929
"github.com/go-logr/logr"
3030
"k8s.io/klog/v2"
3131

32+
"golang.org/x/sync/errgroup"
33+
3234
"github.com/llm-d-incubation/batch-gateway/internal/gc/collector"
3335
gcconfig "github.com/llm-d-incubation/batch-gateway/internal/gc/config"
36+
"github.com/llm-d-incubation/batch-gateway/internal/gc/reconciler"
3437
"github.com/llm-d-incubation/batch-gateway/internal/util/clientset"
3538
ucom "github.com/llm-d-incubation/batch-gateway/internal/util/com"
3639
"github.com/llm-d-incubation/batch-gateway/internal/util/interrupt"
@@ -66,19 +69,35 @@ func run() error {
6669

6770
cfg.DBClientCfg.RedisCfg.ServiceName = "batch-gc"
6871

69-
clients, err := clientset.NewClientset(ctx, ucom.ComponentGC,
72+
clientOpts := []clientset.Option{
7073
clientset.WithDB(cfg.DBClientCfg),
7174
clientset.WithFile(cfg.FileClientCfg),
72-
)
75+
}
76+
if cfg.Reconciler.Enabled {
77+
clientOpts = append(clientOpts, clientset.WithExchange(cfg.DBClientCfg.RedisCfg))
78+
}
79+
80+
clients, err := clientset.NewClientset(ctx, ucom.ComponentGC, clientOpts...)
7381
if err != nil {
7482
return fmt.Errorf("failed to create clients: %w", err)
7583
}
7684
defer func() { _ = clients.Close() }()
7785

7886
gc := collector.NewGarbageCollector(clients.BatchDB, clients.FileDB, clients.File, cfg.DryRun, cfg.Interval, cfg.MaxConcurrency, nil)
7987

80-
if err := gc.RunLoop(ctx); err != nil && ctx.Err() == nil {
81-
return fmt.Errorf("garbage collector failed: %w", err)
88+
g, gCtx := errgroup.WithContext(ctx)
89+
g.Go(func() error { return gc.RunLoop(gCtx) })
90+
91+
if cfg.Reconciler.Enabled {
92+
rec, err := reconciler.NewReconciler(clients.BatchDB, clients.Queue, clients.InFlight, cfg.Reconciler.Interval, cfg.DryRun, nil)
93+
if err != nil {
94+
return fmt.Errorf("failed to create reconciler: %w", err)
95+
}
96+
g.Go(func() error { return rec.RunLoop(gCtx) })
97+
}
98+
99+
if err := g.Wait(); err != nil && ctx.Err() == nil {
100+
return fmt.Errorf("gc/reconciler failed: %w", err)
82101
}
83102

84103
logger.Info("Garbage collector shut down gracefully")

‎internal/database/api/batch_item.go‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -25,6 +25,7 @@ type BatchItem struct {
2525
// BatchQuery specifies parameters for retrieving batches from the database.
2626
type BatchQuery struct {
2727
BaseQuery
28+
NonTerminal bool
2829
}
2930

3031
// BatchDBClient is the typed database client for batch objects.

‎internal/database/mock/mock_inflight_client.go‎

Lines changed: 9 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -78,6 +78,15 @@ func (m *MockInFlightClient) InFlightGetAll(_ context.Context) (map[string]*api.
7878
return result, nil
7979
}
8080

81+
// SetLastSeen overrides the LastSeen timestamp for a specific entry (test helper).
82+
func (m *MockInFlightClient) SetLastSeen(jobID string, lastSeen int64) {
83+
m.mu.Lock()
84+
defer m.mu.Unlock()
85+
if entry, ok := m.entries[jobID]; ok {
86+
entry.LastSeen = lastSeen
87+
}
88+
}
89+
8190
func (m *MockInFlightClient) Close() error {
8291
m.mu.Lock()
8392
defer m.mu.Unlock()

‎internal/database/postgresql/batch_db.go‎

Lines changed: 22 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -20,14 +20,30 @@ import (
2020
"context"
2121
_ "embed"
2222
"fmt"
23+
"strings"
2324

2425
"github.com/go-logr/logr"
26+
2527
"github.com/llm-d-incubation/batch-gateway/internal/database/api"
28+
"github.com/llm-d-incubation/batch-gateway/internal/shared/openai"
2629
)
2730

2831
//go:embed batch_schema.sql
2932
var batchSchemaSql string
3033

34+
// nonTerminalCondition is the SQL WHERE clause fragment that filters for
35+
// non-terminal batch statuses. Computed once since terminal statuses are fixed.
36+
var nonTerminalCondition = buildNonTerminalCondition()
37+
38+
func buildNonTerminalCondition() string {
39+
statuses := openai.TerminalStatuses()
40+
quoted := make([]string, len(statuses))
41+
for i, s := range statuses {
42+
quoted[i] = "'" + string(s) + "'"
43+
}
44+
return colStatus + `::jsonb->>'status' NOT IN (` + strings.Join(quoted, ",") + `)`
45+
}
46+
3147
// Compile-time check: batchDescriptor implements TableDescriptor.
3248
var _ TableDescriptor = (*batchDescriptor)(nil)
3349

@@ -83,8 +99,13 @@ func (c *PostgresBatchDBClient) DBGet(
8399
return
84100
}
85101

102+
var rawConditions []string
103+
if query.NonTerminal {
104+
rawConditions = append(rawConditions, nonTerminalCondition)
105+
}
106+
86107
indexes, contents, _, cursor, expectMore, err := c.get(
87-
ctx, &query.BaseQuery, includeStatic, start, limit, nil)
108+
ctx, &query.BaseQuery, includeStatic, start, limit, nil, rawConditions)
88109
if err != nil {
89110
return
90111
}

‎internal/database/postgresql/batch_db_test.go‎

Lines changed: 61 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -247,6 +247,67 @@ func TestBatchGet(t *testing.T) {
247247
}
248248
})
249249

250+
t.Run("gets non-terminal items", func(t *testing.T) {
251+
client, mock := newTestBatchClient(t)
252+
defer mock.Close()
253+
254+
item := newTestBatchItem("batch-1", testTenantID)
255+
tags, _ := packTags(item.Tags)
256+
257+
rows := pgxmock.NewRows([]string{colID, colTenantID, colExpiry, colTags, colStatus, colSpec}).
258+
AddRow(item.ID, item.TenantID, item.Expiry, &tags, item.Status, item.Spec)
259+
260+
mock.ExpectQuery("SELECT .+ FROM "+testTable+" WHERE .+NOT IN").
261+
WithArgs(0, 11).
262+
WillReturnRows(rows)
263+
264+
items, _, _, err := client.DBGet(ctx,
265+
&api.BatchQuery{NonTerminal: true},
266+
true, 0, 10)
267+
if err != nil {
268+
t.Fatalf("expected no error, got %v", err)
269+
}
270+
if len(items) != 1 {
271+
t.Fatalf("expected 1 item, got %d", len(items))
272+
}
273+
274+
if err := mock.ExpectationsWereMet(); err != nil {
275+
t.Fatalf("unmet expectations: %v", err)
276+
}
277+
})
278+
279+
t.Run("gets non-terminal items filtered by tenant", func(t *testing.T) {
280+
client, mock := newTestBatchClient(t)
281+
defer mock.Close()
282+
283+
item := newTestBatchItem("batch-1", testTenantID)
284+
tags, _ := packTags(item.Tags)
285+
286+
rows := pgxmock.NewRows([]string{colID, colTenantID, colExpiry, colTags, colStatus, colSpec}).
287+
AddRow(item.ID, item.TenantID, item.Expiry, &tags, item.Status, item.Spec)
288+
289+
mock.ExpectQuery("SELECT .+ FROM "+testTable+" WHERE .+NOT IN").
290+
WithArgs(testTenantID, 0, 11).
291+
WillReturnRows(rows)
292+
293+
items, _, _, err := client.DBGet(ctx,
294+
&api.BatchQuery{
295+
BaseQuery: api.BaseQuery{TenantID: testTenantID},
296+
NonTerminal: true,
297+
},
298+
true, 0, 10)
299+
if err != nil {
300+
t.Fatalf("expected no error, got %v", err)
301+
}
302+
if len(items) != 1 {
303+
t.Fatalf("expected 1 item, got %d", len(items))
304+
}
305+
306+
if err := mock.ExpectationsWereMet(); err != nil {
307+
t.Fatalf("unmet expectations: %v", err)
308+
}
309+
})
310+
250311
t.Run("gets items by tag selectors with OR", func(t *testing.T) {
251312
client, mock := newTestBatchClient(t)
252313
defer mock.Close()

‎internal/database/postgresql/db_core.go‎

Lines changed: 7 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -173,9 +173,10 @@ func (c *pgCore) store(ctx context.Context, idx *api.BaseIndexes, contents *api.
173173

174174
// buildGetQuery constructs the SELECT SQL and args from the given query parameters.
175175
// extraFilters is keyed by column name and adds equality conditions for each entry.
176+
// rawConditions are appended verbatim to the WHERE clause (no parameterization).
176177
func (c *pgCore) buildGetQuery(
177178
bq *api.BaseQuery, includeStatic bool, start, limit int,
178-
extraFilters map[string]any,
179+
extraFilters map[string]any, rawConditions []string,
179180
) (string, []any) {
180181
cols := c.selectColumns(includeStatic)
181182

@@ -223,6 +224,8 @@ func (c *pgCore) buildGetQuery(
223224
argIdx++
224225
}
225226

227+
conditions = append(conditions, rawConditions...)
228+
226229
if len(conditions) == 0 {
227230
return "", nil
228231
}
@@ -301,15 +304,16 @@ func (c *pgCore) scanRow(rows pgx.Rows, includeStatic bool) (*api.BaseIndexes, *
301304

302305
// get retrieves items from the database.
303306
// extraFilters adds type-specific equality conditions (e.g., purpose = "batch").
307+
// rawConditions are appended verbatim to the WHERE clause.
304308
func (c *pgCore) get(
305309
ctx context.Context, bq *api.BaseQuery, includeStatic bool, start, limit int,
306-
extraFilters map[string]any,
310+
extraFilters map[string]any, rawConditions []string,
307311
) (
308312
indexes []*api.BaseIndexes, contents []*api.BaseContents, extras []map[string]string,
309313
cursor int, expectMore bool, err error,
310314
) {
311315
// Request one extra row beyond the limit to determine if more results exist.
312-
sql, args := c.buildGetQuery(bq, includeStatic, start, limit+1, extraFilters)
316+
sql, args := c.buildGetQuery(bq, includeStatic, start, limit+1, extraFilters, rawConditions)
313317
if sql == "" {
314318
return nil, nil, nil, 0, false, nil
315319
}

‎internal/database/postgresql/db_core_test.go‎

Lines changed: 5 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -160,7 +160,7 @@ func TestCoreGet_EmptyQuery(t *testing.T) {
160160
defer mock.Close()
161161

162162
indexes, contents, _, cursor, expectMore, err := core.get(
163-
context.Background(), &api.BaseQuery{}, true, 0, 10, nil)
163+
context.Background(), &api.BaseQuery{}, true, 0, 10, nil, nil)
164164
if err != nil {
165165
t.Fatalf("expected no error, got %v", err)
166166
}
@@ -183,7 +183,7 @@ func TestCoreGet_DBFailure(t *testing.T) {
183183
WillReturnError(fmt.Errorf("connection refused"))
184184

185185
_, _, _, _, _, err := core.get(
186-
context.Background(), &api.BaseQuery{TenantID: "t1"}, true, 0, 10, nil)
186+
context.Background(), &api.BaseQuery{TenantID: "t1"}, true, 0, 10, nil, nil)
187187
if err == nil {
188188
t.Fatal("expected error on DB failure")
189189
}
@@ -202,7 +202,7 @@ func TestCoreGet_Expired(t *testing.T) {
202202
WillReturnRows(rows)
203203

204204
indexes, _, _, _, _, err := core.get(
205-
context.Background(), &api.BaseQuery{Expired: true}, true, 0, 10, nil)
205+
context.Background(), &api.BaseQuery{Expired: true}, true, 0, 10, nil, nil)
206206
if err != nil {
207207
t.Fatalf("expected no error, got %v", err)
208208
}
@@ -232,7 +232,7 @@ func TestCoreGet_Pagination(t *testing.T) {
232232
WillReturnRows(rows)
233233

234234
indexes, _, _, cursor, expectMore, err := core.get(
235-
context.Background(), &api.BaseQuery{TenantID: "t1"}, false, 0, 2, nil)
235+
context.Background(), &api.BaseQuery{TenantID: "t1"}, false, 0, 2, nil, nil)
236236
if err != nil {
237237
t.Fatalf("expected no error, got %v", err)
238238
}
@@ -265,7 +265,7 @@ func TestCoreGet_Pagination(t *testing.T) {
265265
WillReturnRows(rows)
266266

267267
indexes, _, _, cursor, expectMore, err := core.get(
268-
context.Background(), &api.BaseQuery{TenantID: "t1"}, false, 0, 2, nil)
268+
context.Background(), &api.BaseQuery{TenantID: "t1"}, false, 0, 2, nil, nil)
269269
if err != nil {
270270
t.Fatalf("expected no error, got %v", err)
271271
}

‎internal/database/postgresql/file_db.go‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -96,7 +96,7 @@ func (c *PostgresFileDBClient) DBGet(
9696
}
9797

9898
indexes, contents, extras, cursor, expectMore, err := c.get(
99-
ctx, &query.BaseQuery, includeStatic, start, limit, extraFilters)
99+
ctx, &query.BaseQuery, includeStatic, start, limit, extraFilters, nil)
100100
if err != nil {
101101
return
102102
}

‎internal/database/redis/redis_db.go‎

Lines changed: 14 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -233,7 +233,7 @@ func (c *DSClientRedis) dBDelete(ctx context.Context, IDs []string, itemType, lo
233233
func (c *DSClientRedis) dbGet(
234234
ctx context.Context, itemType, logPref string, start, limit int, includeStatic bool,
235235
IDs []string, tagSelectors db_api.Tags, tagsLogicalCond db_api.LogicalCond,
236-
expired bool, tenantID, purpose string) (res []any, err error) {
236+
expired bool, tenantID, purpose string, nonTerminal bool) (res []any, err error) {
237237

238238
if ctx == nil {
239239
ctx = context.Background()
@@ -293,6 +293,17 @@ func (c *DSClientRedis) dbGet(
293293
return
294294
}
295295

296+
} else if nonTerminal {
297+
298+
cctx, ccancel := context.WithTimeout(ctx, c.timeout)
299+
defer ccancel()
300+
res, err = redisScriptGetByNonTerminal.Run(cctx, c.redisClient,
301+
[]string{}, getKeyPatternForStore(itemType),
302+
start, limit, tenantID, includeSpec).Slice()
303+
if err != nil {
304+
return
305+
}
306+
296307
} else if len(tenantID) > 0 {
297308

298309
cctx, ccancel := context.WithTimeout(ctx, c.timeout)
@@ -325,7 +336,7 @@ func (c *BatchDBClientRedis) DBGet(
325336

326337
var res []any
327338
res, err = c.dbGet(ctx, itemTypeBatch, "DBGet[Batch]", start, limit, includeStatic,
328-
query.IDs, query.TagSelectors, query.TagsLogicalCond, query.Expired, query.TenantID, "")
339+
query.IDs, query.TagSelectors, query.TagsLogicalCond, query.Expired, query.TenantID, "", query.NonTerminal)
329340
if err != nil {
330341
return
331342
}
@@ -356,7 +367,7 @@ func (c *FileDBClientRedis) DBGet(
356367

357368
var res []any
358369
res, err = c.dbGet(ctx, itemTypeFile, "DBGet[File]", start, limit, includeStatic,
359-
query.IDs, query.TagSelectors, query.TagsLogicalCond, query.Expired, query.TenantID, query.Purpose)
370+
query.IDs, query.TagSelectors, query.TagsLogicalCond, query.Expired, query.TenantID, query.Purpose, false)
360371
if err != nil {
361372
return
362373
}

‎internal/database/redis/redis_ds_client.go‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -88,6 +88,10 @@ var (
8888
getByTenantLua string
8989
redisScriptGetByTenant = goredis.NewScript(commonLua + "\n" + getByTenantLua)
9090

91+
//go:embed redis_get_by_non_terminal.lua
92+
getByNonTerminalLua string
93+
redisScriptGetByNonTerminal = goredis.NewScript(commonLua + "\n" + getByNonTerminalLua)
94+
9195
//go:embed redis_update_cas.lua
9296
updateCASLua string
9397
redisScriptUpdateCAS = goredis.NewScript(updateCASLua)

0 commit comments

Comments
 (0)