wehub-resource-sync
378 lines
14 KiB
Python
378 lines
14 KiB
Python
"""Tests for per-engine thread isolation (issue #1248)."""
|
|
|
|
import concurrent.futures
|
|
from types import SimpleNamespace
|
|
from unittest.mock import MagicMock, patch
|
|
|
|
import mlx.core as mx
|
|
import pytest
|
|
|
|
from omlx.engine_core import EngineCore
|
|
from omlx.request import Request, SamplingParams
|
|
from omlx.scheduler import Scheduler, SchedulerConfig
|
|
|
|
|
|
class TestSchedulerStreamParam:
|
|
"""Scheduler must accept an explicit stream and use it instead of the
|
|
module-level generation_stream."""
|
|
|
|
def test_scheduler_stores_explicit_stream(self):
|
|
mock_model = MagicMock()
|
|
mock_model.model_type = "test"
|
|
mock_tokenizer = MagicMock()
|
|
mock_tokenizer.eos_token_id = 0
|
|
|
|
stream = mx.new_thread_local_stream(mx.default_device())
|
|
scheduler = Scheduler(
|
|
model=mock_model,
|
|
tokenizer=mock_tokenizer,
|
|
stream=stream,
|
|
)
|
|
assert scheduler._stream is stream
|
|
|
|
def test_scheduler_defaults_to_generation_stream(self):
|
|
from omlx.scheduler import _default_generation_stream
|
|
|
|
mock_model = MagicMock()
|
|
mock_model.model_type = "test"
|
|
mock_tokenizer = MagicMock()
|
|
mock_tokenizer.eos_token_id = 0
|
|
|
|
scheduler = Scheduler(
|
|
model=mock_model,
|
|
tokenizer=mock_tokenizer,
|
|
)
|
|
assert scheduler._stream is _default_generation_stream
|
|
|
|
|
|
class TestSchedulerStreamIsolation:
|
|
"""Scheduler must use self._stream in all GPU stream operations,
|
|
never the module-level generation_stream."""
|
|
|
|
def test_no_module_level_generation_stream_in_hot_path(self):
|
|
"""After migration, scheduler.py should not reference the module-level
|
|
generation_stream anywhere in the Scheduler class body except the
|
|
__init__ default fallback and comments/docstrings."""
|
|
import inspect
|
|
import re
|
|
|
|
import omlx.scheduler as sched_mod
|
|
source = inspect.getsource(sched_mod.Scheduler)
|
|
|
|
# Find bare generation_stream references that aren't:
|
|
# - _default_generation_stream (the import alias)
|
|
# - Part of a larger word
|
|
bare_refs = re.findall(
|
|
r'(?<!_default_)(?<!self\._)(?<!\w)generation_stream(?!\w)',
|
|
source,
|
|
)
|
|
# Filter out string literals and comments by checking lines
|
|
code_refs = []
|
|
for line in source.splitlines():
|
|
stripped = line.strip()
|
|
if stripped.startswith('#') or stripped.startswith('"') or stripped.startswith("'"):
|
|
continue
|
|
matches = re.findall(
|
|
r'(?<!_default_)(?<!self\._)(?<!\w)generation_stream(?!\w)',
|
|
line,
|
|
)
|
|
code_refs.extend(matches)
|
|
|
|
assert len(code_refs) == 0, (
|
|
f"Found {len(code_refs)} bare generation_stream references in "
|
|
f"Scheduler class body. All should be self._stream."
|
|
)
|
|
|
|
@pytest.mark.skipif(not mx.metal.is_available(), reason="Metal is required")
|
|
def test_prefill_paths_use_engine_stream_on_worker_thread(self):
|
|
"""Direct prefill paths must bind native lazy ops to the engine stream.
|
|
|
|
Qwen3.5/3.6's native MoE weighted-sum path is activated at 1024
|
|
tokens. Without an explicit stream context, it binds its lazy
|
|
primitive to the executor thread's unrelated default stream and M5
|
|
macOS 27 builds fail during mx.eval (issue #2170).
|
|
"""
|
|
|
|
observed_streams = []
|
|
|
|
class RecordingModel:
|
|
model_type = "test"
|
|
config = SimpleNamespace(model_type="test")
|
|
|
|
def parameters(self):
|
|
return {}
|
|
|
|
def __call__(self, inputs, **kwargs):
|
|
observed_streams.append(mx.default_stream(mx.gpu))
|
|
|
|
class Tokenizer:
|
|
eos_token_id = 2
|
|
pad_token_id = 0
|
|
bos_token_id = 1
|
|
|
|
def encode(self, text, add_special_tokens=True):
|
|
return [1]
|
|
|
|
def decode(self, token_ids, skip_special_tokens=True):
|
|
return ""
|
|
|
|
stream = mx.new_thread_local_stream(mx.default_device())
|
|
scheduler = Scheduler(
|
|
model=RecordingModel(),
|
|
tokenizer=Tokenizer(),
|
|
config=SchedulerConfig(prefill_step_size=2048),
|
|
stream=stream,
|
|
)
|
|
request = Request(
|
|
request_id="external-prefill-stream",
|
|
prompt=[1, 2, 3],
|
|
sampling_params=SamplingParams(),
|
|
)
|
|
request.prompt_token_ids = [1, 2, 3]
|
|
request.num_prompt_tokens = 3
|
|
cache = [SimpleNamespace(state=mx.array([0]))]
|
|
spec_request = Request(
|
|
request_id="specprefill-stream",
|
|
prompt=[1, 2, 3, 4],
|
|
sampling_params=SamplingParams(),
|
|
)
|
|
spec_request.prompt_token_ids = [1, 2, 3, 4]
|
|
spec_request.remaining_tokens = [1, 2, 3, 4]
|
|
spec_request.num_prompt_tokens = 4
|
|
spec_request._specprefill_enabled = True
|
|
spec_request._specprefill_threshold = 1
|
|
spec_request._specprefill_keep_pct = 0.5
|
|
scheduler._specprefill_draft_model = object()
|
|
scheduler._draft_prefix_cache = None
|
|
|
|
def record_specprefill_stream(*args, **kwargs):
|
|
observed_streams.append(mx.default_stream(mx.gpu))
|
|
return mx.ones(4), []
|
|
|
|
def run_prefill():
|
|
worker_default = mx.default_stream(mx.gpu)
|
|
with mx.stream(stream):
|
|
expected_engine_stream = mx.default_stream(mx.gpu)
|
|
scheduler._do_external_prefill(
|
|
request,
|
|
request.prompt_token_ids,
|
|
cache,
|
|
)
|
|
state = scheduler._begin_prefill(
|
|
request,
|
|
request.prompt_token_ids,
|
|
cache,
|
|
)
|
|
scheduler._step_prefill_chunk(state)
|
|
scheduler._try_specprefill_scoring(spec_request)
|
|
return worker_default, expected_engine_stream, mx.default_stream(mx.gpu)
|
|
|
|
with (
|
|
patch(
|
|
"omlx.patches.specprefill.score_tokens",
|
|
side_effect=record_specprefill_stream,
|
|
),
|
|
concurrent.futures.ThreadPoolExecutor(max_workers=1) as executor,
|
|
):
|
|
worker_default, expected_engine_stream, restored_default = executor.submit(
|
|
run_prefill
|
|
).result()
|
|
|
|
assert observed_streams == [
|
|
expected_engine_stream,
|
|
expected_engine_stream,
|
|
expected_engine_stream,
|
|
]
|
|
assert expected_engine_stream != worker_default
|
|
assert restored_default == worker_default
|
|
|
|
|
|
class TestMtpStreamIsolation:
|
|
"""MTP patch must not use the module-level generation_stream directly."""
|
|
|
|
def test_mtp_patch_no_get_generation_stream(self):
|
|
"""_get_generation_stream must not exist — MTP inherits the stream
|
|
from the enclosing BatchGenerator context."""
|
|
import omlx.patches.mlx_lm_mtp.batch_generator as mtp_mod
|
|
|
|
assert not hasattr(mtp_mod, "_get_generation_stream"), (
|
|
"_get_generation_stream still exists in MTP patch; "
|
|
"MTP should inherit the per-engine stream from BatchGenerator"
|
|
)
|
|
|
|
def test_mtp_source_no_module_level_stream_read(self):
|
|
"""MTP patch source must not read sys.modules generation_stream."""
|
|
import inspect
|
|
import omlx.patches.mlx_lm_mtp.batch_generator as mtp_mod
|
|
|
|
source = inspect.getsource(mtp_mod)
|
|
assert "generation_stream" not in source, (
|
|
"MTP patch references generation_stream — all stream context "
|
|
"should be inherited from the enclosing BatchGenerator"
|
|
)
|
|
|
|
|
|
class TestPerEngineExecutor:
|
|
"""Each EngineCore must create its own ThreadPoolExecutor, not share
|
|
a global singleton."""
|
|
|
|
def test_two_engines_have_different_executors(self):
|
|
mock_model_a = MagicMock()
|
|
mock_model_a.model_type = "test"
|
|
mock_model_b = MagicMock()
|
|
mock_model_b.model_type = "test"
|
|
mock_tokenizer = MagicMock()
|
|
mock_tokenizer.eos_token_id = 0
|
|
|
|
with patch("omlx.engine_core.get_registry") as mock_registry:
|
|
mock_registry.return_value.acquire.return_value = True
|
|
|
|
engine_a = EngineCore(mock_model_a, mock_tokenizer)
|
|
engine_b = EngineCore(mock_model_b, mock_tokenizer)
|
|
|
|
assert engine_a._mlx_executor is not engine_b._mlx_executor
|
|
assert engine_a._mlx_stream is not engine_b._mlx_stream
|
|
|
|
engine_a.close()
|
|
engine_b.close()
|
|
|
|
def test_engine_passes_stream_to_scheduler(self):
|
|
mock_model = MagicMock()
|
|
mock_model.model_type = "test"
|
|
mock_tokenizer = MagicMock()
|
|
mock_tokenizer.eos_token_id = 0
|
|
|
|
with patch("omlx.engine_core.get_registry") as mock_registry:
|
|
mock_registry.return_value.acquire.return_value = True
|
|
|
|
engine = EngineCore(mock_model, mock_tokenizer)
|
|
assert engine.scheduler._stream is engine._mlx_stream
|
|
|
|
engine.close()
|
|
|
|
def test_close_clears_compile_cache_then_shuts_down(self):
|
|
"""Normal path (compile-cache clear available): close() clears the
|
|
worker thread's MLX thread_local compile cache (so ~CompilerCache is a
|
|
no-op at thread exit) and then shuts the executor down normally."""
|
|
import omlx.engine_core as ec
|
|
|
|
mock_model = MagicMock()
|
|
mock_model.model_type = "test"
|
|
mock_tokenizer = MagicMock()
|
|
mock_tokenizer.eos_token_id = 0
|
|
|
|
with patch("omlx.engine_core.get_registry") as mock_registry, patch(
|
|
"omlx.engine_core.compile_cache_clear_available", return_value=True
|
|
), patch("omlx.engine_core.clear_thread_compile_cache") as mock_clear, patch(
|
|
"omlx.engine_core._final_engine_thread_reclaim"
|
|
) as mock_reclaim:
|
|
mock_registry.return_value.acquire.return_value = True
|
|
|
|
engine = EngineCore(mock_model, mock_tokenizer)
|
|
executor = engine._mlx_executor
|
|
events = []
|
|
|
|
def reclaim_side_effect(_stream):
|
|
events.append("reclaim")
|
|
assert engine.model is None
|
|
assert engine.tokenizer is None
|
|
assert engine.scheduler is None
|
|
|
|
mock_model.release_resources.side_effect = lambda: events.append("release")
|
|
mock_reclaim.side_effect = reclaim_side_effect
|
|
mock_clear.side_effect = lambda: events.append("compile")
|
|
engine.close()
|
|
|
|
# Memory is reclaimed after dropping engine refs, then the compile
|
|
# cache is cleared on the worker thread before thread shutdown.
|
|
mock_model.release_resources.assert_called_once_with()
|
|
mock_reclaim.assert_called_once_with(engine._mlx_stream)
|
|
mock_clear.assert_called()
|
|
assert events == ["release", "reclaim", "compile"]
|
|
assert engine._mlx_executor is None
|
|
assert executor._shutdown
|
|
assert executor not in ec._immortal_mlx_executors
|
|
|
|
def test_close_keeps_executor_alive_when_clear_unavailable(self):
|
|
"""Fallback (clear symbol unresolvable, e.g. a future MLX rename):
|
|
close() must NOT exit the worker thread, since that would run MLX's
|
|
thread_local ~CompilerCache and crash for @mx.compile models. The
|
|
executor + stream are pinned immortal for the process lifetime."""
|
|
import omlx.engine_core as ec
|
|
|
|
mock_model = MagicMock()
|
|
mock_model.model_type = "test"
|
|
mock_tokenizer = MagicMock()
|
|
mock_tokenizer.eos_token_id = 0
|
|
|
|
with patch("omlx.engine_core.get_registry") as mock_registry, patch(
|
|
"omlx.engine_core.compile_cache_clear_available", return_value=False
|
|
):
|
|
mock_registry.return_value.acquire.return_value = True
|
|
|
|
engine = EngineCore(mock_model, mock_tokenizer)
|
|
executor = engine._mlx_executor
|
|
stream = engine._mlx_stream
|
|
engine.close()
|
|
|
|
assert engine._mlx_executor is None
|
|
assert not executor._shutdown
|
|
assert executor in ec._immortal_mlx_executors
|
|
assert stream in ec._immortal_mlx_streams
|
|
|
|
|
|
class TestConcurrentStreamIsolation:
|
|
"""Verify that per-engine streams don't leak across engines during
|
|
concurrent execution."""
|
|
|
|
def test_concurrent_schedulers_use_own_streams(self):
|
|
"""Two schedulers running step() concurrently must each use their
|
|
own stream, not cross-contaminate."""
|
|
mock_model_a = MagicMock()
|
|
mock_model_a.model_type = "test"
|
|
mock_model_b = MagicMock()
|
|
mock_model_b.model_type = "test"
|
|
mock_tokenizer = MagicMock()
|
|
mock_tokenizer.eos_token_id = 0
|
|
|
|
stream_a = mx.new_thread_local_stream(mx.default_device())
|
|
stream_b = mx.new_thread_local_stream(mx.default_device())
|
|
assert stream_a is not stream_b
|
|
|
|
sched_a = Scheduler(
|
|
model=mock_model_a,
|
|
tokenizer=mock_tokenizer,
|
|
stream=stream_a,
|
|
)
|
|
sched_b = Scheduler(
|
|
model=mock_model_b,
|
|
tokenizer=mock_tokenizer,
|
|
stream=stream_b,
|
|
)
|
|
|
|
assert sched_a._stream is stream_a
|
|
assert sched_b._stream is stream_b
|
|
assert sched_a._stream is not sched_b._stream
|
|
|
|
def test_module_level_generation_stream_unchanged(self):
|
|
"""Creating schedulers with explicit streams must not modify the
|
|
module-level _default_generation_stream."""
|
|
from omlx.scheduler import _default_generation_stream
|
|
|
|
original_id = id(_default_generation_stream)
|
|
stream = mx.new_thread_local_stream(mx.default_device())
|
|
|
|
mock_model = MagicMock()
|
|
mock_model.model_type = "test"
|
|
mock_tokenizer = MagicMock()
|
|
mock_tokenizer.eos_token_id = 0
|
|
|
|
_ = Scheduler(
|
|
model=mock_model,
|
|
tokenizer=mock_tokenizer,
|
|
stream=stream,
|
|
)
|
|
|
|
from omlx.scheduler import _default_generation_stream as current
|
|
assert id(current) == original_id
|