Fix the MLX server dying at startup on CI, and show why it did

src/plamo_translate/servers/mlx/plamo3.py

@@ -284,7 +284,11 @@ def load_model(path: str | Path, dtype: mx.Dtype = mx.float16, lazy: bool = Fals
     weights = {}
     for shard in sorted(path.glob("model*.safetensors")):
         weights.update(mx.load(str(shard)))
-    model.load_weights([(name, array.astype(dtype)) for name, array in weights.items()])
+    # Cast the two largest tensors on the CPU, which gives the same values and has no Metal time limit.
+    on_cpu = {"model.embed_tokens.weight", "lm_head.weight"}
+    model.load_weights(
+        [(name, array.astype(dtype, stream=mx.cpu if name in on_cpu else None)) for name, array in weights.items()]
+    )
     if not lazy:
         mx.eval(model.parameters())
     return model, config

src/plamo_translate/servers/mlx/plamo3_runtime.py

@@ -27,6 +27,8 @@
 from plamo_translate.servers.warnings import suppress_optional_gpu_dependency_warnings
 
 GROUP_SIZE = 64
+# lm_head rows quantized per command buffer: about 70 MB of FP16, a third of a layer.
+LM_HEAD_ROWS = 13440
 # Largest FP32 logit difference accepted between the kernels and the plain MLX ops.
 SELF_CHECK_TOLERANCE = 0.25
 
@@ -57,6 +59,15 @@ def predicate(path: str, module: nn.Module) -> bool | dict:
     nn.quantize(model, group_size=GROUP_SIZE, bits=bits, class_predicate=predicate)
 
 
+def quantize_rows(module: nn.QuantizedLinear, weight: mx.array) -> None:
+    """Quantize an evaluated weight a block of rows at a time; groups never span rows, so the result is the same."""
+    blocks = []
+    for start in range(0, weight.shape[0], LM_HEAD_ROWS):
+        blocks.append(mx.quantize(weight[start : start + LM_HEAD_ROWS], module.group_size, module.bits))
+        mx.eval(blocks[-1])
+    module.weight, module.scales, module.biases = (mx.concatenate(parts, stream=mx.cpu) for parts in zip(*blocks))
+
+
 def checkpoint_path(name: str) -> Path:
     """A local checkpoint directory, or a Hugging Face repository downloaded into the cache."""
     local = Path(name).expanduser()
@@ -82,12 +93,20 @@ def __init__(self, path: str, precision: str | None = None):
         self.precision = precision or "4bit"
         # FP16 keeps more mantissa bits than BF16; activations stayed far below its range on real text.
         self.model, self.config = load_model(self.path, mx.float16, lazy=True)
+        head = self.model.lm_head.weight
         quantize(self.model, self.precision)
         # Convert and quantize one layer at a time, returning freed buffers: in a single evaluation the 4bit load
         # peaked at 7.2 GB and left 5.3 GB cached, too much for 8 GB Macs.
         for layer in self.model.model.layers.layers:
             mx.eval(layer.parameters())
             mx.clear_cache()
+        # The embedding and lm_head are each over twice a layer's size. On CI runners short of memory, Metal aborted
+        # command buffers converting them on the GPU (GPU Timeout Error), so load_model casts both on the CPU and
+        # lm_head is quantized in blocks of rows.
+        mx.eval(self.model.model.embed_tokens.parameters(), head)
+        if isinstance(self.model.lm_head, nn.QuantizedLinear):
+            quantize_rows(self.model.lm_head, head)
+        del head
         mx.eval(self.model.parameters())
         mx.clear_cache()
         with suppress_optional_gpu_dependency_warnings():

tests/test_cli_integration.py

@@ -1,7 +1,9 @@
 import os
+import signal
 import subprocess
 import sys
 import time
+from pathlib import Path
 
 import pytest
 
@@ -21,17 +23,35 @@ def integration_test_environment(monkeypatch, tmp_path):
     monkeypatch.setenv("TMPDIR", str(tmp_path))
 
 
-def wait_for_server_ready(timeout: int = SERVER_STARTUP_TIMEOUT_SECONDS) -> None:
+def start_server(log_path: Path, *arguments: str) -> subprocess.Popen:
+    """Write the server's output to a file: an unread pipe would block the server once it fills."""
+    with open(log_path, "w") as log:
+        return subprocess.Popen(
+            ["plamo-translate", "server", *arguments],
+            stdout=log,
+            stderr=subprocess.STDOUT,
+            # A timed-out server is aborted below: print where each thread was, and keep what it printed before.
+            env={**os.environ, "PYTHONFAULTHANDLER": "1", "PYTHONUNBUFFERED": "1"},
+        )
+
+
+def wait_for_server_ready(
+    process: subprocess.Popen, log_path: Path, timeout: int = SERVER_STARTUP_TIMEOUT_SECONDS
+) -> None:
     deadline = time.monotonic() + timeout
-    while time.monotonic() < deadline:
-        if check_server_running():
-            return
+    while not check_server_running():
+        assert process.poll() is None, f"The server exited with {process.returncode}:\n{log_path.read_text()}"
+        if time.monotonic() >= deadline:
+            process.send_signal(signal.SIGABRT)
+            try:
+                process.wait(timeout=10)
+            except subprocess.TimeoutExpired:
+                pass
+            raise AssertionError(f"Timed out waiting for the MCP server to become ready:\n{log_path.read_text()}")
         time.sleep(0.5)
 
-    raise AssertionError("Timed out waiting for the MCP server to become ready.")
-
 
-def stop_subprocess(process: subprocess.Popen[str] | None) -> None:
+def stop_subprocess(process: subprocess.Popen | None) -> None:
     if process is None:
         return
 
@@ -43,17 +63,13 @@ def stop_subprocess(process: subprocess.Popen[str] | None) -> None:
         process.wait(timeout=5)
 
 
-def test_plamo_translate_server_roundtrip_with_real_model(monkeypatch):
+def test_plamo_translate_server_roundtrip_with_real_model(monkeypatch, tmp_path):
     monkeypatch.delenv("PLAMO_TRANSLATE_CLI_MODEL_NAME", raising=False)
+    log_path = tmp_path / "server.log"
     server_process = None
     try:
-        server_process = subprocess.Popen(
-            ["plamo-translate", "server"],
-            stdout=subprocess.PIPE,
-            stderr=subprocess.PIPE,
-            text=True,
-        )
-        wait_for_server_ready()
+        server_process = start_server(log_path)
+        wait_for_server_ready(server_process, log_path)
 
         config = update_config()
         assert config.get("port") == PLAMO_TRANSLATE_CLI_SERVER_START_PORT
@@ -91,15 +107,11 @@ def test_plamo_translate_4bit_completions_issue_10(monkeypatch, tmp_path):
     monkeypatch.delenv("PLAMO_TRANSLATE_CLI_REPETITION_PENALTY", raising=False)
     monkeypatch.delenv("PLAMO_TRANSLATE_CLI_REPETITION_CONTEXT_SIZE", raising=False)
 
+    log_path = tmp_path / "server.log"
     server_process = None
     try:
-        server_process = subprocess.Popen(
-            ["plamo-translate", "server", "--precision", "4bit"],
-            stdout=subprocess.PIPE,
-            stderr=subprocess.PIPE,
-            text=True,
-        )
-        wait_for_server_ready()
+        server_process = start_server(log_path, "--precision", "4bit")
+        wait_for_server_ready(server_process, log_path)
         assert update_config().get("model_name") == model_name
 
         result = subprocess.run(