diff --git a/mlxrunner/pipeline.go b/mlxrunner/pipeline.go index 7d7b3dd97..390f6e143 100644 --- a/mlxrunner/pipeline.go +++ b/mlxrunner/pipeline.go @@ -90,6 +90,15 @@ func (r *Runner) Prepare(request *Request) (err error) { // The runner serializes requests today so we just use a fixed slot ID. const pipelineSlot = 0 +// skipCleanupOnPanic keeps a backend failure from being replaced by cleanup +// that evaluates or saves partially updated cache state. +func skipCleanupOnPanic(cleanup func()) { + if failure := recover(); failure != nil { + panic(failure) + } + cleanup() +} + func (r *Runner) TextGenerationPipeline(ctx context.Context, request Request) (err error) { mlx.ResetPeakMemory() mlx.Scoped(func() { err = r.generate(ctx, request) }) @@ -106,16 +115,16 @@ func (r *Runner) generate(ctx context.Context, request Request) error { inputs := request.Tokens session := r.cache.begin(inputs, request.MediaItems) - defer session.close() + defer skipCleanupOnPanic(session.close) caches := session.caches media := r.openMedia(request) - defer media.close() + defer skipCleanupOnPanic(media.close) // Built before prefill so a drafter with draft caches follows the prompt // through prefill alongside the target. spec := r.spec.open(request, media.rowLayout()) - defer spec.close() + defer skipCleanupOnPanic(spec.close) seed, position, promptEval, err := r.prefill(ctx, session, spec, media) if err != nil { @@ -137,7 +146,7 @@ func (r *Runner) generate(ctx context.Context, request Request) error { } else { d = r.pipelinedDecoder(nil, caches, seed.ExpandDims(-1), position, media.rowLayout(), grammar) } - defer d.close() + defer skipCleanupOnPanic(d.close) return r.decode(ctx, request, session, d, promptEval) } @@ -253,12 +262,12 @@ type decoder interface { // not streamed or counted. func (r *Runner) decode(ctx context.Context, request Request, session *cacheSession, d decoder, promptEval time.Duration) error { // A sampled-but-undelivered result is still a produced token; record it. - defer func() { + defer skipCleanupOnPanic(func() { results, _, _ := d.drain() for _, res := range results { session.outputs = append(session.outputs, res.Token.Int()) } - }() + }) detok := detokenizer{ tokenizer: r.Tokenizer, diff --git a/mlxrunner/pipeline_panic_test.go b/mlxrunner/pipeline_panic_test.go new file mode 100644 index 000000000..1b0ef0c02 --- /dev/null +++ b/mlxrunner/pipeline_panic_test.go @@ -0,0 +1,99 @@ +package mlxrunner + +import ( + "context" + "testing" + + "github.com/ollama/ollama/api" + "github.com/ollama/ollama/mlx" + "github.com/ollama/ollama/mlx/mlxtest" + "github.com/ollama/ollama/mlxrunner/batch" + "github.com/ollama/ollama/mlxrunner/cache" + sampler "github.com/ollama/ollama/mlxrunner/sample" +) + +type panicDecoder struct { + drained bool +} + +func (*panicDecoder) next(int) ([]sampler.Result, error) { panic("decode failed") } +func (d *panicDecoder) drain() ([]sampler.Result, int, error) { + d.drained = true + panic("second drain") +} +func (*panicDecoder) close() {} + +// In #18856, a second drain replaced the original MLX panic. +func TestDecodePreservesPanic(t *testing.T) { + mlxtest.Run(t, func(t *mlxtest.T) { + d := &panicDecoder{} + request := Request{CompletionRequest: CompletionRequest{Options: api.Options{NumPredict: 1}}} + var failure any + func() { + defer func() { failure = recover() }() + _ = (&Runner{}).decode(context.Background(), request, &cacheSession{inputs: []int32{1}}, d, 0) + }() + if failure != "decode failed" { + t.Fatalf("panic = %v, want original decode failure", failure) + } + if d.drained { + t.Fatal("decoder drained after a panic") + } + }) +} + +type panicForwardModel struct{ textOnlyModel } + +func (panicForwardModel) Forward(_ *batch.Batch, caches []cache.Cache) (*mlx.Array, *mlx.Array) { + caches[0].(*fakeRewindableCache).feed([]int32{1, 2, 3}) + panic("forward failed") +} + +// In #18856, cache close replaced the original panic with an offset error. +func TestGeneratePreservesPanicWithoutSavingCache(t *testing.T) { + mlxtest.Run(t, func(t *mlxtest.T) { + kv := &fakeRewindableCache{tracker: &snapshotTracker{}} + pc := newPrefixCache([]cache.Cache{kv}) + r := &Runner{Model: panicForwardModel{}, cache: pc} + var failure any + func() { + defer func() { failure = recover() }() + mlx.Scoped(func() { + _ = r.generate(context.Background(), Request{Tokens: []int32{1, 2}}) + }) + }() + if failure != "forward failed" { + t.Fatalf("panic = %v, want original forward failure", failure) + } + if len(pc.root.children) != 0 { + t.Fatalf("saved %d cache paths after failed forward", len(pc.root.children)) + } + }) +} + +func TestGenerateStillSavesCacheOnSuccess(t *testing.T) { + mlxtest.Run(t, func(t *mlxtest.T) { + kv := &fakeRewindableCache{tracker: &snapshotTracker{}} + pc := newPrefixCache([]cache.Cache{kv}) + tok := newTestTokenizer(t, []int32{7}) + r := &Runner{ + Model: &fakeMTPModel{predict: map[int32]int32{1: 2, 2: 3, 3: 7}, tok: tok}, + Tokenizer: tok, + Sampler: sampler.New(4096), + cache: pc, + } + request := Request{ + Tokens: []int32{1, 2}, + Responses: make(chan CompletionResponse, 4), + CompletionRequest: CompletionRequest{Options: api.Options{NumPredict: 1}}, + } + mlx.Scoped(func() { + if err := r.generate(context.Background(), request); err != nil { + t.Fatalf("generate: %v", err) + } + }) + if len(pc.root.children) != 1 { + t.Fatalf("saved cache paths = %d, want 1", len(pc.root.children)) + } + }) +}