Skip to content

server: --allowed-models to restrict which model paths a request may load - #1893

Open
mrcushen-arch wants to merge 1 commit into
ml-explore:mainfrom
mrcushen-arch:feat/allowed-models
Open

mrcushen-arch wants to merge 1 commit into
ml-explore:mainfrom
mrcushen-arch:feat/allowed-models

Conversation

@mrcushen-arch

Copy link
Copy Markdown

Summary

Adds --allowed-models to mlx_lm.server: an optional allow-list of model paths a request may load. The --model given at start is always allowed. A request naming anything else is refused before any weights are touched, through the existing error path (HTTP 404 with a JSON error), instead of loading or swapping the resident model. Default behaviour (load on demand) is unchanged.

Fixes #1892.

Why

On a shared machine several scripts point at one server. Each script that names a different model swaps or adds a resident model for everyone else. Measured on an M4 Pro (48 GB): five models cycled through residency, free memory fell to ~6 GB, and a 10-token request went from ~0.5 s to ~40 s / timeouts. With the server restricted to one model: 1 resident, ~18 GB free, 0.56 s.

Changes

  • ModelProvider.__init__: build allowed_models from the flag (+ the startup --model).
  • ModelProvider.load: raise ValueError when the resolved path is not allowed.
  • --allowed-models argparse flag (nargs="*").
  • Tests: TestAllowedModels — refuses a foreign model, allows listed + default + default_model, no list means any. tests/test_server.py passes in full (36).

…load

On a shared machine every script that names a different model swaps or adds a
resident model for everyone else. With --allowed-models the server refuses any
model path not listed (the --model given at start is always allowed) before
touching weights; the request gets the existing 404 error path. Default is
unchanged (load on demand). Fixes ml-explore#1892.
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

mlx_lm.server: option to restrict which models a request may load (allow-list)

1 participant