documentlanguagemodel Public

Watch 0 Fork 0 Star 0
Python · 14788 bytes Raw Blame History
1 """MLX serve launch artifact generation."""
2
3 from __future__ import annotations
4
5 from pathlib import Path
6
7 import pytest
8
9 from dlm.base_models import BASE_MODELS
10 from dlm.export.errors import ExportError, TargetSmokeError
11 from dlm.export.manifest import load_export_manifest
12 from dlm.export.targets.base import TargetResult
13 from dlm.export.targets.mlx_serve import (
14 LAUNCH_SCRIPT_FILENAME,
15 MLX_SERVE_TARGET,
16 _quote_script_arg,
17 _require_prepared_int,
18 _require_prepared_path,
19 _require_prepared_str,
20 _version_from_dir_name,
21 finalize_mlx_serve_export,
22 prepare_mlx_serve_export,
23 )
24 from dlm.store.manifest import Manifest, load_manifest, save_manifest
25 from dlm.store.paths import for_dlm
26
27 _SPEC = BASE_MODELS["smollm2-135m"]
28
29
30 def _write_adapter(path: Path) -> None:
31 path.mkdir(parents=True)
32 (path / "adapter_config.json").write_text("{}", encoding="utf-8")
33 (path / "adapter_model.safetensors").write_bytes(b"adapter")
34
35
36 def _fake_stage_mlx(src: Path, dst: Path, *, base_hf_id: str) -> Path:
37 assert src.exists()
38 assert base_hf_id == _SPEC.hf_id
39 dst.mkdir(parents=True, exist_ok=True)
40 (dst / "adapter_config.json").write_text("{}", encoding="utf-8")
41 (dst / "adapters.safetensors").write_bytes(b"mlx-adapter")
42 return dst
43
44
45 def _setup_flat_store(tmp_path: Path) -> object:
46 store = for_dlm("01MLXTEST", home=tmp_path)
47 store.ensure_layout()
48 save_manifest(store.manifest, Manifest(dlm_id="01MLXTEST", base_model=_SPEC.key))
49 adapter = store.adapter_version(3)
50 _write_adapter(adapter)
51 store.set_current_adapter(adapter)
52 return store
53
54
55 def _setup_named_store(tmp_path: Path) -> object:
56 store = for_dlm("01MLXMULTI", home=tmp_path)
57 store.ensure_layout()
58 save_manifest(store.manifest, Manifest(dlm_id="01MLXMULTI", base_model=_SPEC.key))
59 knowledge = store.adapter_version_for("knowledge", 2)
60 tone = store.adapter_version_for("tone", 4)
61 _write_adapter(knowledge)
62 _write_adapter(tone)
63 store.set_current_adapter_for("knowledge", knowledge)
64 store.set_current_adapter_for("tone", tone)
65 return store
66
67
68 class TestPrepareMlxServeExport:
69 def test_prepare_method_is_not_used_directly(self) -> None:
70 with pytest.raises(NotImplementedError, match="prepare_mlx_serve_export"):
71 MLX_SERVE_TARGET.prepare(object())
72
73 def test_prepare_writes_launch_script_and_manifest(
74 self, tmp_path: Path, monkeypatch: object
75 ) -> None:
76 store = _setup_flat_store(tmp_path)
77 monkeypatch.setattr("dlm.export.targets.mlx_serve.is_apple_silicon", lambda: True)
78 monkeypatch.setattr("dlm.export.targets.mlx_serve.mlx_available", lambda: True)
79 monkeypatch.setattr("dlm.export.targets.mlx_serve.stage_mlx_adapter_dir", _fake_stage_mlx)
80
81 prepared = prepare_mlx_serve_export(
82 store=store,
83 spec=_SPEC,
84 adapter_name=None,
85 adapter_path_override=None,
86 declared_adapter_names=None,
87 )
88 manifest_path = finalize_mlx_serve_export(
89 store=store,
90 spec=_SPEC,
91 prepared=prepared,
92 smoke_output_first_line="hello from mlx",
93 adapter_name=None,
94 adapter_mix=None,
95 )
96
97 assert prepared.launch_script_path is not None
98 assert prepared.launch_script_path.name == LAUNCH_SCRIPT_FILENAME
99 script = prepared.launch_script_path.read_text(encoding="utf-8")
100 assert script.startswith("#!/usr/bin/env bash\nset -euo pipefail\n")
101 assert "python -m mlx_lm.server" in script
102 assert f"--model {_SPEC.hf_id}" in script
103 assert '--adapter-path "$SCRIPT_DIR/adapter"' in script
104
105 export_manifest = load_export_manifest(prepared.export_dir)
106 assert manifest_path == prepared.manifest_path
107 assert export_manifest.target == "mlx-serve"
108 assert export_manifest.quant == "hf"
109 assert export_manifest.adapter_version == 3
110 assert any(artifact.path == "mlx_serve_launch.sh" for artifact in export_manifest.artifacts)
111 assert any(
112 artifact.path == "adapter/adapters.safetensors"
113 for artifact in export_manifest.artifacts
114 )
115
116 store_manifest = load_manifest(store.manifest)
117 assert store_manifest.exports[-1].target == "mlx-serve"
118 assert store_manifest.exports[-1].quant == "hf"
119 assert store_manifest.exports[-1].smoke_output_first_line == "hello from mlx"
120
121 def test_prepare_replaces_stale_staged_adapter_dir(
122 self, tmp_path: Path, monkeypatch: object
123 ) -> None:
124 store = _setup_flat_store(tmp_path)
125 export_dir = store.exports / "mlx-serve"
126 stale_dir = export_dir / "adapter"
127 stale_dir.mkdir(parents=True)
128 (stale_dir / "stale.txt").write_text("stale", encoding="utf-8")
129 monkeypatch.setattr("dlm.export.targets.mlx_serve.is_apple_silicon", lambda: True)
130 monkeypatch.setattr("dlm.export.targets.mlx_serve.mlx_available", lambda: True)
131 monkeypatch.setattr("dlm.export.targets.mlx_serve.stage_mlx_adapter_dir", _fake_stage_mlx)
132
133 prepared = prepare_mlx_serve_export(
134 store=store,
135 spec=_SPEC,
136 adapter_name=None,
137 adapter_path_override=None,
138 declared_adapter_names=None,
139 )
140
141 assert prepared.launch_script_path is not None
142 assert not (prepared.export_dir / "adapter" / "stale.txt").exists()
143 assert (prepared.export_dir / "adapter" / "adapters.safetensors").exists()
144
145 def test_multi_adapter_export_requires_explicit_selection(
146 self, tmp_path: Path, monkeypatch: object
147 ) -> None:
148 store = _setup_named_store(tmp_path)
149 monkeypatch.setattr("dlm.export.targets.mlx_serve.is_apple_silicon", lambda: True)
150 monkeypatch.setattr("dlm.export.targets.mlx_serve.mlx_available", lambda: True)
151
152 with pytest.raises(ExportError, match="one adapter at a time"):
153 prepare_mlx_serve_export(
154 store=store,
155 spec=_SPEC,
156 adapter_name=None,
157 adapter_path_override=None,
158 declared_adapter_names=("knowledge", "tone"),
159 )
160
161 def test_refuses_without_apple_silicon_runtime(
162 self, tmp_path: Path, monkeypatch: object
163 ) -> None:
164 store = _setup_flat_store(tmp_path)
165 monkeypatch.setattr("dlm.export.targets.mlx_serve.is_apple_silicon", lambda: False)
166
167 with pytest.raises(ExportError, match="Apple Silicon"):
168 prepare_mlx_serve_export(
169 store=store,
170 spec=_SPEC,
171 adapter_name=None,
172 adapter_path_override=None,
173 declared_adapter_names=None,
174 )
175
176 def test_refuses_without_mlx_extra(self, tmp_path: Path, monkeypatch: object) -> None:
177 store = _setup_flat_store(tmp_path)
178 monkeypatch.setattr("dlm.export.targets.mlx_serve.is_apple_silicon", lambda: True)
179 monkeypatch.setattr("dlm.export.targets.mlx_serve.mlx_available", lambda: False)
180
181 with pytest.raises(ExportError, match="mlx extra"):
182 prepare_mlx_serve_export(
183 store=store,
184 spec=_SPEC,
185 adapter_name=None,
186 adapter_path_override=None,
187 declared_adapter_names=None,
188 )
189
190 def test_missing_named_adapter_raises(self, tmp_path: Path, monkeypatch: object) -> None:
191 store = _setup_named_store(tmp_path)
192 monkeypatch.setattr("dlm.export.targets.mlx_serve.is_apple_silicon", lambda: True)
193 monkeypatch.setattr("dlm.export.targets.mlx_serve.mlx_available", lambda: True)
194
195 with pytest.raises(ExportError, match="no current adapter under"):
196 prepare_mlx_serve_export(
197 store=store,
198 spec=_SPEC,
199 adapter_name="missing",
200 adapter_path_override=None,
201 declared_adapter_names=None,
202 )
203
204 def test_named_adapter_export_uses_named_dir(self, tmp_path: Path, monkeypatch: object) -> None:
205 store = _setup_named_store(tmp_path)
206 monkeypatch.setattr("dlm.export.targets.mlx_serve.is_apple_silicon", lambda: True)
207 monkeypatch.setattr("dlm.export.targets.mlx_serve.mlx_available", lambda: True)
208 monkeypatch.setattr("dlm.export.targets.mlx_serve.stage_mlx_adapter_dir", _fake_stage_mlx)
209
210 prepared = prepare_mlx_serve_export(
211 store=store,
212 spec=_SPEC,
213 adapter_name="knowledge",
214 adapter_path_override=None,
215 declared_adapter_names=None,
216 )
217
218 assert str(prepared.extras["adapter_dir"]).endswith("knowledge")
219 assert prepared.extras["adapter_version"] == 2
220
221 def test_missing_adapter_override_raises(self, tmp_path: Path, monkeypatch: object) -> None:
222 store = _setup_flat_store(tmp_path)
223 monkeypatch.setattr("dlm.export.targets.mlx_serve.is_apple_silicon", lambda: True)
224 monkeypatch.setattr("dlm.export.targets.mlx_serve.mlx_available", lambda: True)
225
226 with pytest.raises(ExportError, match="adapter_path_override .* does not exist"):
227 prepare_mlx_serve_export(
228 store=store,
229 spec=_SPEC,
230 adapter_name=None,
231 adapter_path_override=tmp_path / "missing",
232 declared_adapter_names=None,
233 )
234
235 def test_existing_adapter_override_uses_mixed_dir(
236 self, tmp_path: Path, monkeypatch: object
237 ) -> None:
238 store = _setup_flat_store(tmp_path)
239 override = tmp_path / "custom-adapter"
240 _write_adapter(override)
241 monkeypatch.setattr("dlm.export.targets.mlx_serve.is_apple_silicon", lambda: True)
242 monkeypatch.setattr("dlm.export.targets.mlx_serve.mlx_available", lambda: True)
243 monkeypatch.setattr("dlm.export.targets.mlx_serve.stage_mlx_adapter_dir", _fake_stage_mlx)
244
245 prepared = prepare_mlx_serve_export(
246 store=store,
247 spec=_SPEC,
248 adapter_name=None,
249 adapter_path_override=override,
250 declared_adapter_names=None,
251 )
252
253 assert str(prepared.extras["adapter_dir"]).endswith("mixed")
254 assert prepared.extras["adapter_version"] == 1
255
256 def test_missing_default_adapter_raises(self, tmp_path: Path, monkeypatch: object) -> None:
257 store = for_dlm("01EMPTYMLX", home=tmp_path)
258 store.ensure_layout()
259 save_manifest(store.manifest, Manifest(dlm_id="01EMPTYMLX", base_model=_SPEC.key))
260 monkeypatch.setattr("dlm.export.targets.mlx_serve.is_apple_silicon", lambda: True)
261 monkeypatch.setattr("dlm.export.targets.mlx_serve.mlx_available", lambda: True)
262
263 with pytest.raises(ExportError, match="no current adapter under"):
264 prepare_mlx_serve_export(
265 store=store,
266 spec=_SPEC,
267 adapter_name=None,
268 adapter_path_override=None,
269 declared_adapter_names=None,
270 )
271
272
273 class TestMlxServeSmoke:
274 def test_smoke_uses_absolute_runtime_paths(self, tmp_path: Path, monkeypatch: object) -> None:
275 store = _setup_flat_store(tmp_path)
276 monkeypatch.setattr("dlm.export.targets.mlx_serve.is_apple_silicon", lambda: True)
277 monkeypatch.setattr("dlm.export.targets.mlx_serve.mlx_available", lambda: True)
278 monkeypatch.setattr("dlm.export.targets.mlx_serve.stage_mlx_adapter_dir", _fake_stage_mlx)
279 prepared = prepare_mlx_serve_export(
280 store=store,
281 spec=_SPEC,
282 adapter_name=None,
283 adapter_path_override=None,
284 declared_adapter_names=None,
285 )
286 seen: list[list[str]] = []
287
288 def _fake_smoke(argv: list[str], **_: object) -> str:
289 seen.append(list(argv))
290 return "mlx replied"
291
292 monkeypatch.setattr("dlm.export.targets.mlx_serve.smoke_openai_compat_server", _fake_smoke)
293
294 result = MLX_SERVE_TARGET.smoke_test(prepared)
295
296 assert result.attempted is True
297 assert result.ok is True
298 assert result.detail == "mlx replied"
299 argv = seen[0]
300 assert argv[:3] == ["python", "-m", "mlx_lm.server"]
301 assert "$SCRIPT_DIR" not in " ".join(argv)
302 assert _SPEC.hf_id in argv
303 assert str(prepared.export_dir / "adapter") in argv
304
305 def test_smoke_failure_returns_failed_result(self, tmp_path: Path, monkeypatch: object) -> None:
306 store = _setup_flat_store(tmp_path)
307 monkeypatch.setattr("dlm.export.targets.mlx_serve.is_apple_silicon", lambda: True)
308 monkeypatch.setattr("dlm.export.targets.mlx_serve.mlx_available", lambda: True)
309 monkeypatch.setattr("dlm.export.targets.mlx_serve.stage_mlx_adapter_dir", _fake_stage_mlx)
310 prepared = prepare_mlx_serve_export(
311 store=store,
312 spec=_SPEC,
313 adapter_name=None,
314 adapter_path_override=None,
315 declared_adapter_names=None,
316 )
317
318 def _fake_smoke(argv: list[str], **_: object) -> str:
319 _ = argv
320 raise TargetSmokeError("boom")
321
322 monkeypatch.setattr("dlm.export.targets.mlx_serve.smoke_openai_compat_server", _fake_smoke)
323
324 result = MLX_SERVE_TARGET.smoke_test(prepared)
325
326 assert result.attempted is True
327 assert result.ok is False
328 assert result.detail == "boom"
329
330
331 class TestMlxServeHelpers:
332 def test_quote_script_arg_and_int_validation(self) -> None:
333 assert _quote_script_arg("$SCRIPT_DIR/adapter") == '"$SCRIPT_DIR/adapter"'
334 assert _quote_script_arg("plain value") == "'plain value'"
335
336 prepared = TargetResult(
337 name="mlx-serve",
338 export_dir=Path("/tmp/export"),
339 manifest_path=Path("/tmp/export/export_manifest.json"),
340 extras={"adapter_version": "bad"},
341 )
342 with pytest.raises(ExportError, match="missing int extra"):
343 _require_prepared_int(prepared, "adapter_version")
344
345 def test_string_and_path_validation(self) -> None:
346 prepared = TargetResult(
347 name="mlx-serve",
348 export_dir=Path("/tmp/export"),
349 manifest_path=Path("/tmp/export/export_manifest.json"),
350 extras={"model": "", "adapter_dir": "bad"},
351 )
352
353 with pytest.raises(ExportError, match="missing string extra"):
354 _require_prepared_str(prepared, "model")
355 with pytest.raises(ExportError, match="missing Path extra"):
356 _require_prepared_path(prepared, "adapter_dir")
357
358 def test_version_from_dir_name_defaults_for_non_version_dirs(self) -> None:
359 assert _version_from_dir_name(Path("custom-adapter")) == 1