Skip to content

tests: make the MLX boundary guard catch mlx_lm imports - #108

Open
VKVstudio wants to merge 1 commit into
Edge0-AI:mainfrom
VKVstudio:fix/hygiene-mlx-lm-regex
Open

VKVstudio wants to merge 1 commit into
Edge0-AI:mainfrom
VKVstudio:fix/hygiene-mlx-lm-regex

Conversation

@VKVstudio

Copy link
Copy Markdown

What

tests/test_repo_hygiene.py::test_mlx_imports_stay_inside_the_backend enforces the rule stated in backends/__init__.py: MLX packages are imported only under edge0/backends/mlx/. Its regex requires a space or a dot right after mlx, so from mlx_lm.models... import ... never matches and the guard silently passes framework code that imports mlx_lm directly.

This PR tightens the pattern to ^\s*(?:from|import)\s+mlx(?:_lm)?\b.

What it finds

The stricter guard immediately flags src/edge0/streaming/layer.py:49:

from mlx_lm.models.switch_layers import _gather_sort, _scatter_unsort

Rather than turn CI red on an existing file, that path is listed in an explicit _MLX_IMPORT_KNOWN_VIOLATIONS set with a note on what should move into the framework. Remove the entry once the two helpers no longer come from mlx_lm. I'm happy to send that follow-up: they are a few lines of argsort/gather that can live next to moe/routing.py, which would also make the streaming layer importable on a non-MLX backend.

If you would rather have the guard fail loudly instead, drop the allowlist — the regex change alone is the fix.

Checks

  • pytest tests/test_repo_hygiene.py — 4 passed. The file imports nothing from edge0 or mlx, so this runs on any platform, same as the existing Linux hygiene job.
  • The tightened regex matches import mlx, import mlx.core as mx, from mlx_lm import x, from mlx_lm.models.base import ..., and does not match from edge0.backends.mlx._impl import ... or a hypothetical mlxfoo package.

Context

I'm working on a CUDA/PyTorch backend for the reserved backends/cuda/ slot (https://github.com/smartbulldog-pro/edge0-cuda/tree/cuda-backend). This import is the one thing that keeps streaming/layer.py from importing on a non-MLX backend, which is how I ran into the guard.

The regex in test_mlx_imports_stay_inside_the_backend required a space
or a dot right after "mlx", so "from mlx_lm.models... import ..." never
matched and the guard silently passed framework code that imports
mlx_lm directly.

Tighten the pattern to "(from|import)\s+mlx(?:_lm)?\b".  The stricter
guard immediately flags src/edge0/streaming/layer.py:49, which imports
_gather_sort / _scatter_unsort from mlx_lm.models.switch_layers.  Rather
than turn CI red, that file is listed in an explicit
_MLX_IMPORT_KNOWN_VIOLATIONS set with a note on what should move into
the framework; remove the entry once the helpers no longer come from
mlx_lm.

Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
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.

1 participant