Update MLX and MLX-C with threading fixes (#15845)
* Update MLX and MLX-C * Run MLX CGO work on a locked OS thread MLX now relies on OS-thread-local execution state for streams, encoders, and caches. Add an mlxthread executor backed by runtime.LockOSThread and route runner initialization, model load, inference, status memory reads, and cleanup through the worker so Go goroutine migration cannot split MLX state across native threads. Also stop caching default MLX streams before the runner owns the thread and add worker/threaded MLX regression tests. * mlx: use common status writer * mlx: bundle missing libjaccl on arm64 Inspired by #15793 * review comments
This commit is contained in:
@@ -65,7 +65,12 @@ func (w *StatusWriter) AppendError(msg string) {
|
||||
// logs, add a small rolling buffer here to capture those fragments.
|
||||
|
||||
var errorPrefixes = []string{
|
||||
"mlx:",
|
||||
"MLX:",
|
||||
"panic:",
|
||||
"fatal error:",
|
||||
"error:",
|
||||
"Error:",
|
||||
"CUDA error",
|
||||
"ROCm error",
|
||||
"cudaMalloc failed",
|
||||
@@ -79,15 +84,21 @@ var errorPrefixes = []string{
|
||||
|
||||
func (w *StatusWriter) Write(b []byte) (int, error) {
|
||||
var errMsg string
|
||||
errStart := -1
|
||||
var errPrefix string
|
||||
for _, prefix := range errorPrefixes {
|
||||
if _, after, ok := bytes.Cut(b, []byte(prefix)); ok {
|
||||
line := after
|
||||
if j := bytes.IndexByte(line, '\n'); j >= 0 {
|
||||
line = line[:j]
|
||||
}
|
||||
errMsg = prefix + string(bytes.TrimRight(line, " \t\r"))
|
||||
if i := bytes.Index(b, []byte(prefix)); i >= 0 && (errStart < 0 || i < errStart) {
|
||||
errStart = i
|
||||
errPrefix = prefix
|
||||
}
|
||||
}
|
||||
if errStart >= 0 {
|
||||
line := b[errStart+len(errPrefix):]
|
||||
if j := bytes.IndexByte(line, '\n'); j >= 0 {
|
||||
line = line[:j]
|
||||
}
|
||||
errMsg = errPrefix + string(bytes.TrimRight(line, " \t\r"))
|
||||
}
|
||||
if errMsg != "" {
|
||||
w.AppendError(errMsg)
|
||||
}
|
||||
|
||||
@@ -1,35 +1,59 @@
|
||||
package llm
|
||||
|
||||
import (
|
||||
"os"
|
||||
"io"
|
||||
"testing"
|
||||
)
|
||||
|
||||
func TestStatusWriterCapturesErrorLine(t *testing.T) {
|
||||
f, err := os.CreateTemp(t.TempDir(), "status-writer")
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
defer f.Close()
|
||||
|
||||
w := NewStatusWriter(f)
|
||||
if _, err := w.Write([]byte("llama_init_from_model: failed to initialize the context: failed to initialize Metal backend\n")); err != nil {
|
||||
t.Fatal(err)
|
||||
tests := []struct {
|
||||
name string
|
||||
log string
|
||||
want string
|
||||
}{
|
||||
{
|
||||
name: "llama init",
|
||||
log: "llama_init_from_model: failed to initialize the context: failed to initialize Metal backend\n",
|
||||
want: "llama_init_from_model: failed to initialize the context: failed to initialize Metal backend",
|
||||
},
|
||||
{
|
||||
name: "cobra error",
|
||||
log: "Error: foo baz bar\n",
|
||||
want: "Error: foo baz bar",
|
||||
},
|
||||
{
|
||||
name: "uppercase mlx",
|
||||
log: "MLX: there was an error\n",
|
||||
want: "MLX: there was an error",
|
||||
},
|
||||
{
|
||||
name: "panic header",
|
||||
log: "time=2026-05-01T15:36:45.053Z level=INFO source=pipeline.go:71 msg=\"peak memory\" size=\"8.26 GiB\"\n" +
|
||||
"panic: mlx: Failed to compile kernel: nvrtc: error: invalid value for --gpu-architecture (-arch)\n" +
|
||||
"\t. at /go/src/github.com/ollama/ollama/build/_deps/mlx-c-src/mlx/c/transforms.cpp:15\n\n" +
|
||||
"goroutine 31 [running]:\n" +
|
||||
"golang.org/x/sync/errgroup.(*Group).Go.func1()\n" +
|
||||
"\tgolang.org/x/sync@v0.17.0/errgroup/errgroup.go:93 +0x50\n",
|
||||
want: "panic: mlx: Failed to compile kernel: nvrtc: error: invalid value for --gpu-architecture (-arch)",
|
||||
},
|
||||
}
|
||||
|
||||
if got, want := w.LastError(), "llama_init_from_model: failed to initialize the context: failed to initialize Metal backend"; got != want {
|
||||
t.Fatalf("LastError = %q, want %q", got, want)
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
w := NewStatusWriter(io.Discard)
|
||||
if _, err := w.Write([]byte(tt.log)); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
|
||||
if got := w.LastError(); got != tt.want {
|
||||
t.Fatalf("LastError = %q, want %q", got, tt.want)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestStatusWriterAccumulatesErrorLines(t *testing.T) {
|
||||
f, err := os.CreateTemp(t.TempDir(), "status-writer")
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
defer f.Close()
|
||||
|
||||
w := NewStatusWriter(f)
|
||||
w := NewStatusWriter(io.Discard)
|
||||
if _, err := w.Write([]byte("error: failed to initialize the Metal library\n")); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user