mirror of
https://github.com/ollama/ollama.git
synced 2026-09-21 13:38:14 -05:00
Jesse Gross
13037ecb14
The bindings freed arrays by sweeping everything not pinned, so freeing anything required knowing what every other caller still held, and code that never swept accumulated until memory ran out. The prefix cache's eviction of a long stored path did exactly that: each merge copied the KV snapshots and nothing freed the consumed copies until the request ended, which drove a second long request past physical memory. Every array now belongs to a scope. A function scope, entered with Scoped or one of the ScopedEval forms, frees what was created in it when the function returns; results leave only by being returned. A held scope is closed by its holder and frees what was attached to it. A graph is built in a function scope and evaluated after it, so the eval frees each intermediate as it consumes it. Pin, Unpin, Sweep, and the array list's mutex are gone. On an M5 Max with qwen3.8:27b-mlx, the second 84k-token request after a stored one peaks at 35 GB instead of 57 GB; the cold path is unchanged. The copies themselves are untouched, so restoring an owned path can still exceed memory.
72 lines
1.6 KiB
Go
72 lines
1.6 KiB
Go
package create
|
|
|
|
import (
|
|
"fmt"
|
|
"runtime"
|
|
"sync"
|
|
"sync/atomic"
|
|
|
|
"github.com/ollama/ollama/x/mlxrunner/mlx"
|
|
)
|
|
|
|
var (
|
|
mlxThreadOnce sync.Once
|
|
mlxThreadStarted atomic.Bool
|
|
mlxWork chan func()
|
|
mlxInitErr error
|
|
)
|
|
|
|
// runOnMLXThread runs f on the MLX thread and returns its error. The thread is
|
|
// started (and MLX initialized) on first use. A panic in f is recovered and
|
|
// returned as an error so a kernel failure cannot kill the pinned thread.
|
|
//
|
|
// TODO(pdevine): This method should be revisited when the `ollama create` is
|
|
// instead run on the ollama server process instead of the client.
|
|
func runOnMLXThread(f func() error) error {
|
|
mlxThreadOnce.Do(func() {
|
|
mlxWork = make(chan func())
|
|
ready := make(chan error)
|
|
go func() {
|
|
runtime.LockOSThread() // pinned for the process lifetime; never unlocked
|
|
err := mlx.CheckInit()
|
|
if err == nil && mlx.GPUIsAvailable() {
|
|
mlx.SetDefaultDeviceGPU()
|
|
}
|
|
ready <- err
|
|
if err != nil {
|
|
return
|
|
}
|
|
for work := range mlxWork {
|
|
work()
|
|
}
|
|
}()
|
|
mlxInitErr = <-ready
|
|
mlxThreadStarted.Store(mlxInitErr == nil)
|
|
})
|
|
if mlxInitErr != nil {
|
|
return fmt.Errorf("MLX init failed: %w", mlxInitErr)
|
|
}
|
|
|
|
done := make(chan error, 1)
|
|
mlxWork <- func() {
|
|
defer func() {
|
|
if r := recover(); r != nil {
|
|
done <- fmt.Errorf("mlx: %v", r)
|
|
}
|
|
}()
|
|
done <- f()
|
|
}
|
|
return <-done
|
|
}
|
|
|
|
// releaseMLXCache releases the MLX buffer cache. It is a no-op if no MLX work has run.
|
|
func releaseMLXCache() {
|
|
if !mlxThreadStarted.Load() {
|
|
return
|
|
}
|
|
_ = runOnMLXThread(func() error {
|
|
mlx.ClearCache()
|
|
return nil
|
|
})
|
|
}
|