wehub-resource-sync
46 lines
1.6 KiB
Python
46 lines
1.6 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
"""Tests for omlx.utils.compile_cache (thread-local MLX compile-cache clear)."""
|
|
|
|
import concurrent.futures as cf
|
|
|
|
import mlx.core as mx
|
|
|
|
import omlx.utils.compile_cache as cc
|
|
|
|
|
|
def test_clear_available_returns_bool():
|
|
# Contract: never raises, always returns a bool. On a real Apple-Silicon
|
|
# MLX install the symbol resolves, but the helper must degrade gracefully.
|
|
assert isinstance(cc.compile_cache_clear_available(), bool)
|
|
|
|
|
|
def test_clear_thread_compile_cache_does_not_raise():
|
|
# Safe to call on any thread, even with an empty cache.
|
|
cc.clear_thread_compile_cache()
|
|
|
|
|
|
def test_clear_after_compile_on_worker_then_shutdown():
|
|
"""Core scenario behind the fix: run an @mx.compile fn on a worker thread,
|
|
clear that thread's cache on the same thread, then shut the worker down.
|
|
Must not crash (~CompilerCache runs on an empty cache)."""
|
|
|
|
@mx.compile
|
|
def f(x):
|
|
return x * 2 + 1
|
|
|
|
ex = cf.ThreadPoolExecutor(max_workers=1)
|
|
try:
|
|
ex.submit(lambda: mx.eval(f(mx.arange(8)))).result()
|
|
ex.submit(cc.clear_thread_compile_cache).result()
|
|
finally:
|
|
ex.shutdown(wait=True)
|
|
|
|
|
|
def test_noop_when_symbol_unavailable(monkeypatch):
|
|
"""When the symbol cannot be resolved, available() is False and clear() is
|
|
a no-op (callers then fall back to keeping the worker thread alive)."""
|
|
monkeypatch.setattr(cc, "_resolved", True)
|
|
monkeypatch.setattr(cc, "_clear_fn", None)
|
|
assert cc.compile_cache_clear_available() is False
|
|
cc.clear_thread_compile_cache() # must not raise
|