ceos/unsloth
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

feat(cli): support MLX distributed inference (#6845)

Browse source 
* feat(cli): detect MLX distributed launch context

* feat(mlx): wire distributed inference backend

* feat(cli): broadcast MLX distributed chat turns

* fix(cli): wait indefinitely for distributed chat turns

* fix(cli): report MLX distributed load errors cleanly

* fix(mlx): route distributed vlm through loader

* fix(cli): detect inline MLX host JSON

* fix(studio): harden distributed object sharing

* fix(studio): select JACCL distributed backend

* fix(cli): abort distributed error paths

* Distinguish real stream errors from model text via GenStreamError in distributed CLI

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Fail loud when MLX distributed init returns a singleton group

The worker only reaches this block when distributed was explicitly
requested. A singleton (size 1) group means the launch failed to form a
real group (MLX built without distributed support, or an invalid launch
env/hostfile); silently continuing leaves nonzero ranks looping forever
on share_distributed_object. Raise instead so the surrounding handler
returns a clear load error.

* Tighten MLX distributed inference comments

---------

Co-authored-by: Daniel Han <danielhanchen@gmail.com>
Co-authored-by: danielhanchen <unslothai@gmail.com>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
This commit is contained in:
Long Yixing 2026-07-08 18:25:39 +08:00 • committed by GitHub
commit 2a6abe2ff5
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
9 changed files with 1190 additions and 157 deletions
Download patch file Download diff file Expand all files Collapse all files

217
unsloth_cli/_inference.py
View file

@ -4,9 +4,11 @@
"""Model loading and streaming shared by `inference` and `chat`."""
import asyncio
import json
import os
import re
import sys
from contextlib import contextmanager, redirect_stderr, redirect_stdout
from pathlib import Path
from typing import List, Optional
@ -14,10 +16,18 @@ import typer
_THINK_OPEN = "<think>"
_THINK_BLOCK = re.compile(rf"{re.escape(_THINK_OPEN)}.*?</think>", re.DOTALL)
_STREAMED_ERROR_PREFIX = "Error: "
# Cloudflare (in front of remote Studio proxies like RunPod) 403s the default
# "Python-urllib/X.Y" User-Agent as a bot; send a real one on every request.
_USER_AGENT = "unsloth-cli"
_MPI_ENV_PAIRS = (
("OMPI_COMM_WORLD_RANK", "OMPI_COMM_WORLD_SIZE"),
("PMI_RANK", "PMI_SIZE"),
("PMIX_RANK", "PMIX_SIZE"),
("MPI_RANK", "MPI_WORLD_SIZE"),
("MV2_COMM_WORLD_RANK", "MV2_COMM_WORLD_SIZE"),
)
# Built lazily; urllib stays function-local to match this module.
_no_redirect_opener = None
@ -61,6 +71,108 @@ def configure_quiet_logging() -> None:
os.environ.setdefault("HF_HUB_DISABLE_PROGRESS_BARS", "1")
def _parse_nonnegative_int(value: Optional[str]) -> Optional[int]:
if value is None:
return None
try:
parsed = int(value)
except (TypeError, ValueError):
return None
return parsed if parsed >= 0 else None
def _first_mpi_env_pair() -> tuple[Optional[int], Optional[int]]:
for rank_name, size_name in _MPI_ENV_PAIRS:
rank = _parse_nonnegative_int(os.environ.get(rank_name))
world_size = _parse_nonnegative_int(os.environ.get(size_name))
if rank is not None and world_size is not None and world_size > 1 and rank < world_size:
return rank, world_size
return None, None
def _json_rank_count_from_env(name: str) -> Optional[int]:
value = os.environ.get(name)
if not value:
return None
try:
if value.lstrip().startswith(("[", "{")):
data = json.loads(value)
else:
with open(value, "r") as f:
data = json.load(f)
except (OSError, json.JSONDecodeError):
return None
if isinstance(data, list):
return len(data)
if isinstance(data, dict) and isinstance(data.get("hosts"), list):
return len(data["hosts"])
return None
def mlx_distributed_info() -> tuple[bool, int, Optional[int]]:
"""Return launch-context metadata without initializing MLX distributed."""
rank = _parse_nonnegative_int(os.environ.get("MLX_RANK"))
world_size = _parse_nonnegative_int(os.environ.get("MLX_WORLD_SIZE"))
if rank is not None:
if (
world_size is not None
and world_size > 1
and rank < world_size
and os.environ.get("NCCL_HOST_IP")
and os.environ.get("NCCL_PORT")
):
return True, rank, world_size
inferred_size = _json_rank_count_from_env("MLX_HOSTFILE")
if inferred_size is not None and inferred_size > 1 and rank < inferred_size:
return True, rank, inferred_size
inferred_size = _json_rank_count_from_env("MLX_IBV_DEVICES")
if (
inferred_size is not None
and inferred_size > 1
and rank < inferred_size
and os.environ.get("MLX_JACCL_COORDINATOR")
):
return True, rank, inferred_size
return False, 0, None
mpi_rank, mpi_world_size = _first_mpi_env_pair()
return mpi_rank is not None, mpi_rank or 0, mpi_world_size
def mlx_distributed_uses_mpi() -> bool:
"""Whether the current distributed context was launched through MPI."""
return (
_parse_nonnegative_int(os.environ.get("MLX_RANK")) is None
and _first_mpi_env_pair()[0] is not None
)
@contextmanager
def quiet_if_nonzero_mlx_rank():
"""Silence parent and child-process stdout/stderr on nonzero ranks."""
if mlx_distributed_info()[1] == 0:
yield
return
sys.stdout.flush()
sys.stderr.flush()
saved_stdout_fd = os.dup(1)
saved_stderr_fd = os.dup(2)
with open(os.devnull, "w") as devnull:
try:
os.dup2(devnull.fileno(), 1)
os.dup2(devnull.fileno(), 2)
with redirect_stdout(devnull), redirect_stderr(devnull):
yield
finally:
sys.stdout.flush()
sys.stderr.flush()
os.dup2(saved_stdout_fd, 1)
os.dup2(saved_stderr_fd, 2)
os.close(saved_stdout_fd)
os.close(saved_stderr_fd)
def visible_text(text: str, show_thinking: bool) -> str:
if show_thinking:
return text
@ -120,6 +232,21 @@ def collect_stream(stream, show_thinking: bool) -> str:
return visible_text(raw, show_thinking)
def raise_on_streamed_error(stream):
# Match real backend errors by type (GenStreamError), not the "Error:" text
# prefix, so a completion whose text opens with "Error:" is not misread as a
# failure that aborts a distributed run.
try:
ensure_studio_backend_path()
from core.inference.orchestrator import GenStreamError
except Exception:
GenStreamError = None
for chunk in stream:
if GenStreamError is not None and isinstance(chunk, GenStreamError):
raise RuntimeError(str(chunk)[len(_STREAMED_ERROR_PREFIX) :].strip() or "Unknown error")
yield chunk
def render_columns(
left_label: str,
left_text: str,
@ -200,6 +327,19 @@ class ChatBackend:
except Exception:
pass
def share_distributed_object(
self,
obj,
*,
timeout = 300.0,
):
if self._kind != "unsloth" or not hasattr(self._backend, "share_distributed_object"):
raise RuntimeError(
"Distributed MLX chat requires the Unsloth MLX backend; "
f"backend '{self._kind}' cannot broadcast chat turns."
)
return self._backend.share_distributed_object(obj, timeout = timeout)
def resolve_model_config(model: str, *, hf_token: Optional[str]):
ensure_studio_backend_path()
@ -293,36 +433,59 @@ def load_chat_backend(
fresh_backend uses a private orchestrator so a second model (compare's
base column) can run alongside the main one.
"""
if model_config is None:
model_config = resolve_model_config(model, hf_token = hf_token)
with quiet_if_nonzero_mlx_rank():
is_mlx_distributed, rank, _world_size = mlx_distributed_info()
if model_config is None:
model_config = resolve_model_config(model, hf_token = hf_token)
typer.echo(f"Loading {model}", err = True)
if is_mlx_distributed and model_config.is_gguf:
if rank == 0:
typer.echo(
"Distributed MLX inference does not support GGUF/llama.cpp models. "
"Use a non-GGUF MLX model under mlx.launch, or run GGUF without "
"mlx.launch.",
err = True,
)
raise typer.Exit(code = 1)
if model_config.is_gguf:
return _load_gguf_backend(
model_config,
hf_token = hf_token,
max_seq_length = max_seq_length,
tensor_parallel = tensor_parallel,
llama_extra_args = llama_extra_args,
)
if rank == 0:
typer.echo(f"Loading {model}", err = True)
if fresh_backend:
ensure_studio_backend_path()
from core.inference import InferenceOrchestrator
backend = InferenceOrchestrator()
else:
ensure_studio_backend_path()
from core.inference import get_inference_backend
backend = get_inference_backend()
if not backend.load_model(
config = model_config,
max_seq_length = max_seq_length,
load_in_4bit = load_in_4bit,
hf_token = hf_token,
):
typer.echo("Model load failed", err = True)
raise typer.Exit(code = 1)
if model_config.is_gguf:
return _load_gguf_backend(
model_config,
hf_token = hf_token,
max_seq_length = max_seq_length,
tensor_parallel = tensor_parallel,
llama_extra_args = llama_extra_args,
)
if fresh_backend:
ensure_studio_backend_path()
from core.inference import InferenceOrchestrator
backend = InferenceOrchestrator()
else:
ensure_studio_backend_path()
from core.inference import get_inference_backend
backend = get_inference_backend()
try:
loaded = backend.load_model(
config = model_config,
max_seq_length = max_seq_length,
load_in_4bit = load_in_4bit,
hf_token = hf_token,
tensor_parallel = tensor_parallel,
mlx_distributed = is_mlx_distributed,
)
except Exception as exc:
if not is_mlx_distributed:
raise
if rank == 0:
typer.echo(str(exc) or "Model load failed", err = True)
raise typer.Exit(code = 1)
if not loaded:
typer.echo("Model load failed", err = True)
raise typer.Exit(code = 1)
return ChatBackend("unsloth", backend)

178
unsloth_cli/commands/chat.py
View file

@ -1,6 +1,7 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
import sys
from typing import List, Optional
import typer
@ -12,6 +13,10 @@ from unsloth_cli._inference import (
connect_studio_server,
ensure_studio_backend_path,
load_chat_backend,
mlx_distributed_info,
mlx_distributed_uses_mpi,
quiet_if_nonzero_mlx_rank,
raise_on_streamed_error,
render_columns,
resolve_model_config,
stream_markdown,
@ -107,6 +112,20 @@ def _compare_needs_second_model() -> bool:
return False
def _drain_available_stdin() -> None:
"""Drain already-buffered launcher stdin on nonzero distributed ranks."""
try:
import os
from select import select
fd = sys.stdin.fileno()
while select([fd], [], [], 0)[0]:
if not os.read(fd, 8192):
break
except Exception:
return
def _pick_trained_model(console) -> str:
ensure_studio_backend_path()
from utils.models import scan_trained_models
@ -158,7 +177,8 @@ def chat(
"--tensor-parallel/--no-tensor-parallel",
help = (
"Split a GGUF across GPUs by tensor (--split-mode tensor) instead "
"of by layer. Ignored for non-GGUF models."
"of by layer. Under non-MPI mlx.launch, select MLX tensor "
"parallel mode instead of pipeline mode."
),
),
llama_extra_args: Optional[List[str]] = typer.Option(
@ -195,15 +215,43 @@ def chat(
console = Console()
err = Console(stderr = True)
is_mlx_distributed, rank, _world_size = mlx_distributed_info()
should_print = rank == 0
if is_mlx_distributed and mlx_distributed_uses_mpi():
if should_print:
err.print(
"Distributed `unsloth chat` with MPI needs rank-0 prompt broadcast, "
"which is not enabled yet. Use a non-MPI MLX launcher backend "
"such as ring/JACCL for now.",
style = "red",
markup = False,
)
raise typer.Exit(code = 1)
if model is None:
if is_mlx_distributed:
if should_print:
err.print(
"Distributed `unsloth chat` requires an explicit model id or path.",
style = "red",
markup = False,
)
raise typer.Exit(code = 1)
model = _pick_trained_model(console)
# Resolve first so --compare can be rejected before the slow load.
model_config = resolve_model_config(model, hf_token = hf_token)
with quiet_if_nonzero_mlx_rank():
model_config = resolve_model_config(model, hf_token = hf_token)
compare_blocked = _compare_blocked_reason(model_config)
if is_mlx_distributed:
compare_blocked = (
"distributed MLX chat does not support compare mode yet because it "
"would need a second distributed worker group on the same ranks"
)
if compare and compare_blocked:
err.print(f"--compare unavailable: {compare_blocked}", style = "red", markup = False)
if should_print:
err.print(f"--compare unavailable: {compare_blocked}", style = "red", markup = False)
raise typer.Exit(code = 1)
load_opts = dict(
@ -215,9 +263,11 @@ def chat(
)
# Prefer a running Studio server: instant starts, model shared with the UI.
chat_backend = None if no_server else connect_studio_server(model, **load_opts)
chat_backend = (
None if (no_server or is_mlx_distributed) else connect_studio_server(model, **load_opts)
)
server_mode = chat_backend is not None
if server_mode:
if server_mode and should_print:
console.print(
"(Studio server connected — model stays warm after /exit)",
style = "bright_black",
@ -242,23 +292,26 @@ def chat(
return True
base_id = model_config.base_model
if not base_id:
console.print(
"(compare unavailable: this adapter doesn't record its base model)",
style = "yellow",
)
if should_print:
console.print(
"(compare unavailable: this adapter doesn't record its base model)",
style = "yellow",
)
return False
console.print(
f"(loading base model {base_id} for compare — keeps two models in memory)",
style = "bright_black",
markup = False,
)
if should_print:
console.print(
f"(loading base model {base_id} for compare — keeps two models in memory)",
style = "bright_black",
markup = False,
)
try:
# Use the same precision as the tuned model for fair comparison
base_load_opts = dict(load_opts) # Copy original options
base_load_opts["load_in_4bit"] = _get_base_load_in_4bit(model_config)
base_backend = load_chat_backend(base_id, fresh_backend = True, **base_load_opts)
except Exception as exc:
err.print(f"(base model load failed: {exc})", style = "red", markup = False)
if should_print:
err.print(f"(base model load failed: {exc})", style = "red", markup = False)
return False
return True
@ -267,7 +320,7 @@ def chat(
def generate(backend = None, use_adapter = None):
# Reads messages and show_thinking live, so /reset and /think apply.
return (backend or chat_backend).stream(
stream = (backend or chat_backend).stream(
messages,
system_prompt = system_prompt,
temperature = temperature,
@ -278,22 +331,48 @@ def chat(
enable_thinking = show_thinking,
use_adapter = use_adapter,
)
return raise_on_streamed_error(stream) if is_mlx_distributed else stream
console.print()
console.print(f"Chatting with {name}", style = "bold green", markup = False)
console.print(_HELP, style = "bright_black")
if should_print:
console.print()
console.print(f"Chatting with {name}", style = "bold green", markup = False)
console.print(_HELP, style = "bright_black")
# legacy_windows: pre-VT consoles print raw ANSI as ←[1;36m garbage.
you_prompt = _you_prompt(console.is_terminal and not console.legacy_windows)
you_prompt = (
_you_prompt(console.is_terminal and not console.legacy_windows) if should_print else ""
)
assistant_label = "[bold magenta]Assistant:[/bold magenta]"
try:
while True:
try:
user = input(you_prompt).strip()
except (EOFError, KeyboardInterrupt):
console.print()
break
if should_print:
try:
user = input(you_prompt).strip()
except (EOFError, KeyboardInterrupt):
if should_print:
console.print()
user = "/exit"
turn = {"type": "turn", "text": user}
else:
turn = None
if is_mlx_distributed:
try:
turn = chat_backend.share_distributed_object(turn, timeout = None)
if not should_print:
_drain_available_stdin()
except Exception as exc:
if should_print:
err.print(
f"\n(error sharing chat turn: {exc})",
style = "red",
markup = False,
)
raise typer.Exit(code = 1)
if not turn:
continue
user = str(turn.get("text", "")).strip()
if not user:
continue
@ -301,55 +380,69 @@ def chat(
break
if user == "/reset":
messages = []
console.print("(history cleared)", style = "bright_black")
if should_print:
console.print("(history cleared)", style = "bright_black")
continue
if user == "/think":
show_thinking = not show_thinking
state = "on" if show_thinking else "off"
console.print(f"(thinking {state})", style = "bright_black")
if should_print:
state = "on" if show_thinking else "off"
console.print(f"(thinking {state})", style = "bright_black")
continue
if user == "/compare":
if compare_blocked:
console.print(f"(compare unavailable: {compare_blocked})", style = "yellow")
if should_print:
console.print(f"(compare unavailable: {compare_blocked})", style = "yellow")
continue
if not compare_mode and dual_compare and not load_base_for_compare():
continue
compare_mode = not compare_mode
state = "on" if compare_mode else "off"
console.print(f"(compare {state})", style = "bright_black")
if should_print:
state = "on" if compare_mode else "off"
console.print(f"(compare {state})", style = "bright_black")
continue
if user in ("/help", "/?"):
console.print(_HELP, style = "bright_black")
if should_print:
console.print(_HELP, style = "bright_black")
continue
messages.append({"role": "user", "content": user})
try:
if compare_mode:
console.print("(comparing base vs tuned…)", style = "bright_black")
if should_print:
console.print("(comparing base vs tuned…)", style = "bright_black")
if dual_compare:
base_text = collect_stream(generate(backend = base_backend), show_thinking)
tuned_text = collect_stream(generate(), show_thinking)
else:
base_text = collect_stream(generate(use_adapter = False), show_thinking)
tuned_text = collect_stream(generate(use_adapter = True), show_thinking)
console.print()
render_columns(
"base", base_text, f"{name} (tuned)", tuned_text, console = console
)
if should_print:
console.print()
render_columns(
"base", base_text, f"{name} (tuned)", tuned_text, console = console
)
# History continues as the tuned model; base is just the reference.
answer = tuned_text
else:
console.print(assistant_label)
answer = stream_markdown(generate(), show_thinking, console = console)
if should_print:
console.print(assistant_label)
answer = stream_markdown(generate(), show_thinking, console = console)
else:
answer = collect_stream(generate(), show_thinking)
except KeyboardInterrupt:
# Ctrl-C aborts this answer only; drop the unanswered turn.
console.print("\n(interrupted)", style = "bright_black")
if should_print:
console.print("\n(interrupted)", style = "bright_black")
messages.pop()
continue
except Exception as exc:
err.print(f"\n(error: {exc})", style = "red", markup = False)
if should_print:
err.print(f"\n(error: {exc})", style = "red", markup = False)
messages.pop()
if is_mlx_distributed:
raise typer.Exit(code = 1)
continue
messages.append(
@ -359,4 +452,5 @@ def chat(
chat_backend.close()
if base_backend is not None:
base_backend.close()
err.print("\nBye.", style = "bright_black")
if should_print:
err.print("\nBye.", style = "bright_black")

47
unsloth_cli/commands/inference.py
View file

@ -6,9 +6,13 @@ from typing import List, Optional
import typer
from unsloth_cli._inference import (
collect_stream,
configure_quiet_logging,
connect_studio_server,
load_chat_backend,
mlx_distributed_info,
mlx_distributed_uses_mpi,
raise_on_streamed_error,
stream_to_stdout,
)
@ -36,7 +40,8 @@ def inference(
"--tensor-parallel/--no-tensor-parallel",
help = (
"Split a GGUF across GPUs by tensor (--split-mode tensor) instead "
"of by layer. Ignored for non-GGUF models."
"of by layer. Under non-MPI mlx.launch, select MLX tensor "
"parallel mode instead of pipeline mode."
),
),
llama_extra_args: Optional[List[str]] = typer.Option(
@ -69,8 +74,20 @@ def inference(
if not verbose:
configure_quiet_logging()
# A running Studio server keeps the model warm between runs, which is
# exactly what a one-shot command wants.
is_mlx_distributed, rank, _world_size = mlx_distributed_info()
if is_mlx_distributed and mlx_distributed_uses_mpi():
if rank == 0:
typer.echo(
"Distributed `unsloth inference` with MPI is not supported by "
"the current subprocess backend. Use a non-MPI MLX launcher "
"backend such as ring/JACCL for now.",
err = True,
)
raise typer.Exit(code = 1)
# A running Studio server keeps the model warm between runs. Under
# mlx.launch, every rank must enter the local MLX path instead of rank 0
# alone talking to a server.
load_opts = dict(
hf_token = hf_token,
max_seq_length = max_seq_length,
@ -78,7 +95,9 @@ def inference(
tensor_parallel = tensor_parallel,
llama_extra_args = llama_extra_args,
)
chat_backend = None if no_server else connect_studio_server(model, **load_opts)
chat_backend = (
None if (no_server or is_mlx_distributed) else connect_studio_server(model, **load_opts)
)
if chat_backend is None:
chat_backend = load_chat_backend(model, **load_opts)
try:
@ -92,7 +111,23 @@ def inference(
repetition_penalty = repetition_penalty,
enable_thinking = think,
)
typer.echo("Assistant:")
stream_to_stdout(stream, show_thinking = think)
if is_mlx_distributed:
stream = raise_on_streamed_error(stream)
if rank == 0:
typer.echo("Assistant:")
try:
stream_to_stdout(stream, show_thinking = think)
except RuntimeError as exc:
if not is_mlx_distributed:
raise
typer.echo(f"Error: {exc}", err = True)
raise typer.Exit(code = 1)
else:
try:
collect_stream(stream, show_thinking = think)
except RuntimeError:
if not is_mlx_distributed:
raise
raise typer.Exit(code = 1)
finally:
chat_backend.close()

330
unsloth_cli/tests/test_inference_chat.py
View file

@ -26,6 +26,8 @@ from unsloth_cli._inference import (
ChatBackend,
HttpChatBackend,
collect_stream,
mlx_distributed_info,
mlx_distributed_uses_mpi,
render_columns,
visible_text,
)
@ -39,6 +41,16 @@ class _FakeConfig:
path = None
_EXPECTED_MPI_ENV_PAIRS = [
("OMPI_COMM_WORLD_RANK", "OMPI_COMM_WORLD_SIZE"),
("PMI_RANK", "PMI_SIZE"),
("PMIX_RANK", "PMIX_SIZE"),
("MPI_RANK", "MPI_WORLD_SIZE"),
("MV2_COMM_WORLD_RANK", "MV2_COMM_WORLD_SIZE"),
]
_IGNORED_DISTRIBUTED_ENV_PAIRS = [("SLURM_PROCID", "SLURM_NTASKS")]
def _chat_app():
cli = typer.Typer()
cli.command()(chatmod.chat)
@ -53,6 +65,39 @@ def _inference_app():
return cli
def _clear_mlx_distributed_env(monkeypatch):
for name in (
"MLX_RANK",
"MLX_HOSTFILE",
"MLX_WORLD_SIZE",
"MLX_IBV_DEVICES",
"MLX_JACCL_COORDINATOR",
"NCCL_HOST_IP",
"NCCL_PORT",
*(rank for rank, _size in _EXPECTED_MPI_ENV_PAIRS + _IGNORED_DISTRIBUTED_ENV_PAIRS),
*(size for _rank, size in _EXPECTED_MPI_ENV_PAIRS + _IGNORED_DISTRIBUTED_ENV_PAIRS),
):
monkeypatch.delenv(name, raising = False)
def _set_mlx_nccl_env(
monkeypatch,
*,
rank: str = "0",
size: str = "2",
):
monkeypatch.setenv("MLX_RANK", rank)
monkeypatch.setenv("MLX_WORLD_SIZE", size)
monkeypatch.setenv("NCCL_HOST_IP", "127.0.0.1")
monkeypatch.setenv("NCCL_PORT", "12345")
@pytest.fixture(autouse = True)
def _isolate_mlx_distributed_env(monkeypatch):
_clear_mlx_distributed_env(monkeypatch)
monkeypatch.delenv("HF_TOKEN", raising = False)
def test_visible_text_passthrough_when_shown():
text = "<think>reasoning</think>answer"
assert visible_text(text, show_thinking = True) == text
@ -101,6 +146,46 @@ def test_inference_exposes_gguf_runtime_options():
assert "--llama-extra-arg" in (getattr(extra, "param_decls", None) or [])
def test_mlx_distributed_info_reads_launch_env(monkeypatch, tmp_path):
_clear_mlx_distributed_env(monkeypatch)
assert mlx_distributed_info() == (False, 0, None)
assert mlx_distributed_uses_mpi() is False
monkeypatch.setenv("MLX_RANK", "1")
monkeypatch.setenv("MLX_WORLD_SIZE", "2")
assert mlx_distributed_info() == (False, 0, None)
monkeypatch.setenv("NCCL_HOST_IP", "127.0.0.1")
monkeypatch.setenv("NCCL_PORT", "12345")
assert mlx_distributed_info() == (True, 1, 2)
assert mlx_distributed_uses_mpi() is False
_clear_mlx_distributed_env(monkeypatch)
ring_hostfile = tmp_path / "ring.json"
ring_hostfile.write_text('[["127.0.0.1:5000"], ["127.0.0.1:5001"]]\n')
monkeypatch.setenv("MLX_RANK", "0")
monkeypatch.setenv("MLX_HOSTFILE", str(ring_hostfile))
assert mlx_distributed_info() == (True, 0, 2)
assert mlx_distributed_uses_mpi() is False
_clear_mlx_distributed_env(monkeypatch)
monkeypatch.setenv("MLX_RANK", "1")
monkeypatch.setenv("MLX_IBV_DEVICES", '[["node-a"], ["node-b"]]')
monkeypatch.setenv("MLX_JACCL_COORDINATOR", "node-a:12345")
assert mlx_distributed_info() == (True, 1, 2)
assert mlx_distributed_uses_mpi() is False
_clear_mlx_distributed_env(monkeypatch)
monkeypatch.setenv("OMPI_COMM_WORLD_RANK", "1")
monkeypatch.setenv("OMPI_COMM_WORLD_SIZE", "2")
assert mlx_distributed_info() == (True, 1, 2)
assert mlx_distributed_uses_mpi() is True
_clear_mlx_distributed_env(monkeypatch)
monkeypatch.setenv("MLX_RANK", "bad")
monkeypatch.setenv("MLX_WORLD_SIZE", "-3")
assert mlx_distributed_info() == (False, 0, None)
def test_chat_command_is_registered_with_options():
params = inspect.signature(chatmod.chat).parameters
assert "model" in params
@ -751,7 +836,6 @@ def test_chat_server_mode_compare_loads_base_locally(monkeypatch):
assert result.exit_code == 0, result.output
assert "(compare on)" in result.output
# Only the base model loaded locally, on its own private backend.
assert base_loads == [("fake/base", True)]
assert streamed == ["base", "tuned"]
assert set(closed) == {"http", "base"}
@ -785,6 +869,248 @@ def test_chat_compare_on_mlx_loads_base_model_side_by_side(monkeypatch):
assert result.exit_code == 0, result.output
assert loads == [("tuned-run", False), ("fake/base", True)]
# Both models answered the turn, via plain generation (no adapter toggle).
assert ("base", None) in streamed and ("tuned", None) in streamed
assert set(closed) == {"tuned", "base"}
@pytest.mark.parametrize(
("chunk_kind", "expected_exit"),
[
("answer", 0),
("model_text_error", 0),
("real_error", 1),
],
)
def test_inference_under_mlx_launch_handles_stream(monkeypatch, chunk_kind, expected_exit):
from unsloth_cli.commands import inference as infermod
from unsloth_cli._inference import ensure_studio_backend_path
ensure_studio_backend_path()
from core.inference.orchestrator import GenStreamError
if chunk_kind == "answer":
chunks = ["answer"]
elif chunk_kind == "model_text_error":
# Model output whose visible text starts with "Error:" must not abort.
chunks = ["Error: printed by the model, not a backend failure"]
else:
chunks = [GenStreamError("Error: generation failed")]
loads, closed = [], []
class _FakeBackend:
def stream(self, messages, **kwargs):
return iter(chunks)
def close(self):
closed.append(True)
_set_mlx_nccl_env(monkeypatch, rank = "0")
monkeypatch.setattr(
infermod,
"connect_studio_server",
lambda *_a, **_k: (_ for _ in ()).throw(AssertionError("server disabled")),
)
monkeypatch.setattr(
infermod,
"load_chat_backend",
lambda model, **kwargs: (loads.append((model, kwargs)), _FakeBackend())[1],
)
result = CliRunner().invoke(
_inference_app(),
["fake-model", "hello", "--tensor-parallel"],
)
assert result.exit_code == expected_exit, result.output
assert loads[0][1]["tensor_parallel"] is True
if chunk_kind == "real_error":
assert "generation failed" in result.output
def test_chat_under_mlx_launch_nonzero_rank_drains_stdin(monkeypatch):
drains, closed = [], []
turns = iter(
[
{"type": "turn", "text": "hi"},
{"type": "turn", "text": "/exit"},
]
)
class _FakeChatBackend:
def share_distributed_object(
self,
obj,
*,
timeout = 300.0,
):
assert obj is None
return next(turns)
def stream(self, messages, **kwargs):
return iter(["hidden"])
def close(self):
closed.append(True)
_set_mlx_nccl_env(monkeypatch, rank = "1")
monkeypatch.setattr(chatmod, "resolve_model_config", lambda *a, **k: _FakeConfig())
monkeypatch.setattr(
chatmod,
"connect_studio_server",
lambda *_a, **_k: (_ for _ in ()).throw(AssertionError("server disabled")),
)
monkeypatch.setattr(chatmod, "load_chat_backend", lambda *a, **k: _FakeChatBackend())
monkeypatch.setattr(chatmod, "_compare_needs_second_model", lambda: False)
monkeypatch.setattr(chatmod, "_drain_available_stdin", lambda: drains.append(True))
result = CliRunner().invoke(_chat_app(), ["fake-model"], input = "hi\n/exit\n")
assert result.exit_code == 0, result.output
assert "Chatting with" not in result.output
assert drains == [True, True]
assert closed == [True]
def test_chat_under_mlx_launch_rank0_bypasses_studio_and_prints(monkeypatch):
loads, shares, closed = [], [], []
class _FakeChatBackend:
def share_distributed_object(
self,
obj,
*,
timeout = 300.0,
):
shares.append((obj, timeout))
return obj
def stream(self, messages, **kwargs):
return iter(["hello"])
def close(self):
closed.append(True)
_set_mlx_nccl_env(monkeypatch, rank = "0")
monkeypatch.setattr(chatmod, "resolve_model_config", lambda *a, **k: _FakeConfig())
monkeypatch.setattr(
chatmod,
"connect_studio_server",
lambda *_a, **_k: (_ for _ in ()).throw(AssertionError("server disabled")),
)
monkeypatch.setattr(
chatmod,
"load_chat_backend",
lambda model, **kwargs: (loads.append((model, kwargs)), _FakeChatBackend())[1],
)
monkeypatch.setattr(chatmod, "_compare_needs_second_model", lambda: False)
result = CliRunner().invoke(
_chat_app(),
["fake-model", "--tensor-parallel"],
input = "hi\n/exit\n",
)
assert result.exit_code == 0, result.output
assert "Chatting with fake-model" in result.output
assert "hello" in result.output
assert loads and loads[0][0] == "fake-model"
assert loads[0][1]["tensor_parallel"] is True
assert shares == [
({"type": "turn", "text": "hi"}, None),
({"type": "turn", "text": "/exit"}, None),
]
@pytest.mark.parametrize(
("stream_error", "expected_exit"),
[("exception", 1), ("chunk", 1), ("model_text", 0)],
)
def test_chat_under_mlx_launch_exits_on_generation_error(monkeypatch, stream_error, expected_exit):
from unsloth_cli._inference import ensure_studio_backend_path
ensure_studio_backend_path()
from core.inference.orchestrator import GenStreamError
closed = []
class _FakeChatBackend:
def share_distributed_object(
self,
obj,
*,
timeout = 300.0,
):
return obj
def stream(self, messages, **kwargs):
if stream_error == "exception":
raise RuntimeError("generation failed")
if stream_error == "model_text":
# Plain model text starting with "Error:" must not abort the run.
return iter(["Error: printed by the model"])
return iter([GenStreamError("Error: generation failed")])
def close(self):
closed.append(True)
_set_mlx_nccl_env(monkeypatch, rank = "0")
monkeypatch.setattr(chatmod, "resolve_model_config", lambda *a, **k: _FakeConfig())
monkeypatch.setattr(
chatmod,
"connect_studio_server",
lambda *_a, **_k: (_ for _ in ()).throw(AssertionError("server disabled")),
)
monkeypatch.setattr(chatmod, "load_chat_backend", lambda *a, **k: _FakeChatBackend())
monkeypatch.setattr(chatmod, "_compare_needs_second_model", lambda: False)
result = CliRunner().invoke(_chat_app(), ["fake-model"], input = "hi\n/exit\n")
assert result.exit_code == expected_exit
if expected_exit:
assert "generation failed" in result.output
assert closed == [True]
def test_load_chat_backend_forwards_mlx_distributed_options(monkeypatch):
import unsloth_cli._inference as inference
calls = []
class _FakeBackend:
def load_model(self, **kwargs):
calls.append(kwargs)
return True
class _FakeModelConfig:
is_gguf = False
@classmethod
def from_identifier(cls, **_kwargs):
return cls()
fake_backend = _FakeBackend()
fake_inference = types.ModuleType("core.inference")
fake_inference.get_inference_backend = lambda: fake_backend
fake_utils = types.ModuleType("utils")
fake_utils.__path__ = []
fake_models = types.ModuleType("utils.models")
fake_models.ModelConfig = _FakeModelConfig
_set_mlx_nccl_env(monkeypatch, rank = "0")
monkeypatch.setitem(sys.modules, "core", types.ModuleType("core"))
monkeypatch.setitem(sys.modules, "core.inference", fake_inference)
monkeypatch.setitem(sys.modules, "utils", fake_utils)
monkeypatch.setitem(sys.modules, "utils.models", fake_models)
monkeypatch.setattr(inference, "ensure_studio_backend_path", lambda: None)
inference.load_chat_backend(
"fake-model",
hf_token = None,
max_seq_length = 2048,
load_in_4bit = True,
tensor_parallel = True,
)
assert calls[0]["tensor_parallel"] is True
assert calls[0]["mlx_distributed"] is True