mlxrunner

Documentation

Index

• func Execute(args []string) error
• type Client
  • func NewClient(modelName string, softContextLength int) (*Client, error)
  • func (c *Client) ApplyChatTemplate(ctx context.Context, req llm.ChatRequest) (string, error)
  • func (c *Client) Chat(ctx context.Context, req llm.ChatRequest, fn func(llm.ChatResponse)) error
  • func (c *Client) Close() error
  • func (c *Client) Completion(ctx context.Context, req llm.CompletionRequest, ...) error
  • func (c *Client) ContextLength() int
  • func (c *Client) Detokenize(ctx context.Context, tokens []int) (string, error)
  • func (c *Client) Embedding(ctx context.Context, input string) ([]float32, int, error)
  • func (c *Client) GetDeviceInfos(ctx context.Context) []ml.DeviceInfo
  • func (c *Client) GetPort() int
  • func (c *Client) HasExited() bool
  • func (c *Client) Load(ctx context.Context, _ ml.SystemInfo, gpus []ml.DeviceInfo, requireFull bool) ([]ml.DeviceID, error)
  • func (c *Client) MemorySize() (total, vram uint64)
  • func (c *Client) ModelPath() string
  • func (c *Client) Pid() int
  • func (c *Client) Ping(ctx context.Context) error
  • func (c *Client) Tokenize(ctx context.Context, content string) ([]int, error)
  • func (c *Client) VRAMByGPU(id ml.DeviceID) uint64
  • func (c *Client) WaitUntilRunning(ctx context.Context) error
• type CompletionRequest
• type CompletionResponse
• type Request
• type Runner
  • func (r *Runner) Load(modelName string) error
  • func (r *Runner) Prepare(request *Request) error
  • func (r *Runner) Run(host, port string, mux http.Handler) error
  • func (r *Runner) TextGenerationPipeline(ctx context.Context, request Request) error

Functions

func Execute

func Execute(args []string) error

Types

type Client

type Client struct {
	// contains filtered or unexported fields
}

Client wraps an MLX runner subprocess to implement llm.LlamaServer for LLM models.

func NewClient

func NewClient(modelName string, softContextLength int) (*Client, error)

NewClient prepares a new MLX runner client for LLM models. The subprocess is not started until Load() is called.

func (*Client) ApplyChatTemplate

func (c *Client) ApplyChatTemplate(ctx context.Context, req llm.ChatRequest) (string, error)

func (*Client) Chat

func (c *Client) Chat(ctx context.Context, req llm.ChatRequest, fn func(llm.ChatResponse)) error

func (*Client) Close

func (c *Client) Close() error

Close terminates the subprocess.

func (*Client) Completion

func (c *Client) Completion(ctx context.Context, req llm.CompletionRequest, fn func(llm.CompletionResponse)) error

Completion implements llm.LlamaServer.

func (*Client) ContextLength

func (c *Client) ContextLength() int

func (*Client) Detokenize

func (c *Client) Detokenize(ctx context.Context, tokens []int) (string, error)

Detokenize implements llm.LlamaServer.

func (*Client) Embedding

func (c *Client) Embedding(ctx context.Context, input string) ([]float32, int, error)

Embedding implements llm.LlamaServer.

func (*Client) GetDeviceInfos

func (c *Client) GetDeviceInfos(ctx context.Context) []ml.DeviceInfo

GetDeviceInfos implements llm.LlamaServer.

func (*Client) GetPort

func (c *Client) GetPort() int

GetPort implements llm.LlamaServer.

func (*Client) HasExited

func (c *Client) HasExited() bool

HasExited implements llm.LlamaServer.

func (*Client) Load

func (c *Client) Load(ctx context.Context, _ ml.SystemInfo, gpus []ml.DeviceInfo, requireFull bool) ([]ml.DeviceID, error)

Load checks whether the model fits in GPU memory and starts the subprocess.

func (*Client) MemorySize

func (c *Client) MemorySize() (total, vram uint64)

MemorySize implements llm.LlamaServer.

func (*Client) ModelPath

func (c *Client) ModelPath() string

ModelPath implements llm.LlamaServer.

func (*Client) Pid

func (c *Client) Pid() int

Pid implements llm.LlamaServer.

func (*Client) Ping

func (c *Client) Ping(ctx context.Context) error

Ping implements llm.LlamaServer.

func (*Client) Tokenize

func (c *Client) Tokenize(ctx context.Context, content string) ([]int, error)

Tokenize implements llm.LlamaServer.

func (*Client) VRAMByGPU

func (c *Client) VRAMByGPU(id ml.DeviceID) uint64

VRAMByGPU implements llm.LlamaServer.

func (*Client) WaitUntilRunning

func (c *Client) WaitUntilRunning(ctx context.Context) error

WaitUntilRunning waits for the subprocess to be ready.

type CompletionRequest

type CompletionRequest struct {
	Prompt      string
	Media       []llm.MediaData
	Options     api.Options
	Logprobs    bool
	TopLogprobs int
}

type CompletionResponse

type CompletionResponse struct {
	Content    string
	Done       bool
	DoneReason int

	PromptEvalCount    int
	PromptEvalDuration time.Duration
	EvalCount          int
	EvalDuration       time.Duration

	Logprobs []llm.Logprob

	Error *api.StatusError
}

type Request

type Request struct {
	CompletionRequest
	Responses chan CompletionResponse
	Pipeline  func(context.Context, Request) error

	Ctx         context.Context //nolint:containedctx // Queued requests carry caller cancellation to the runner.
	Tokens      []int32
	MediaItems  []mediaItem
	Layout      any // opaque PrepareMedia layout state, stamped on every batch
	SamplerOpts sample.Options
}

Request is a short-lived struct that carries a completion request through a channel from the HTTP handler to the runner goroutine. The ctx field must travel with the request so that cancellation propagates across the channel boundary.

type Runner

type Runner struct {
	Model     base.Model
	Tokenizer *tokenizer.Tokenizer
	Requests  chan Request
	Sampler   *sample.Sampler
	// contains filtered or unexported fields
}

func (*Runner) Load

func (r *Runner) Load(modelName string) error

func (*Runner) Prepare

func (r *Runner) Prepare(request *Request) error

Prepare tokenizes the prompt and validates it against the model's context length. It is safe to call from any goroutine. On success it populates request.Tokens and adjusts request.Options.NumPredict.

func (*Runner) Run

func (r *Runner) Run(host, port string, mux http.Handler) error

func (*Runner) TextGenerationPipeline

func (r *Runner) TextGenerationPipeline(ctx context.Context, request Request) error