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

studio: fix double-install of tilelang on the FLA hook install path

Browse source 
Backend CI surfaced a test-isolation bug introduced by the
post_available_fn mechanism for finding #7. The wrapper ran
`post_available_fn` in BOTH paths (install ran AND gate already True),
but `_fla_install` already chains tilelang on the install path, so the
post-available step then called tilelang install AGAIN.

This was masked locally because tilelang was installed in the
workspace venv (post_available short-circuited on
`_tilelang_importable()` returning True). CI starts with no tilelang,
so the second call actually fired and the mock recorded two calls.

Fix: only run `post_available_fn` when the install path did NOT run.
That preserves the finding #7 semantics (tilelang repair when FLA
already True but tilelang missing or tvm-ffi broken) without
duplicating the chained install on the gate-was-False path.

Also tightened `test_hook_skips_install_when_gate_already_true` to
monkeypatch `_tilelang_importable=True` and
`_installed_tvm_ffi_version=0.1.9` so it stays a pure "no install at
all" test regardless of the venv's actual state.
This commit is contained in:
danielhanchen 2026-05-17 01:25:25 +00:00
commit 85cdafc184
2 changed files with 19 additions and 4 deletions
Download patch file Download diff file Expand all files Collapse all files

14
studio/backend/core/training/worker.py
View file

@ -888,7 +888,9 @@ def _install_fast_path_hooks(event_queue: Any, model_name: str) -> None:
except AttributeError:
pass
ok = original()
ran_install = False
if not ok:
ran_install = True
logger.info("Hook fired for %s; triggering install", gate_name)
_send_status(
event_queue,
@ -909,10 +911,14 @@ def _install_fast_path_hooks(event_queue: Any, model_name: str) -> None:
gate_name,
ok,
)
# Even when FLA was already True, the post-available action
# may still have work (tilelang missing / broken tvm-ffi
# repair).
if ok and post_available_fn is not None:
# post_available_fn handles edge cases that ONLY occur on
# the gate-was-already-True path (e.g. tilelang missing
# while FLA is already importable, or apache-tvm-ffi on
# the broken-versions list while FLA otherwise works).
# If install_fn ran, it already chained the matching
# follow-up install (`_fla_install` installs tilelang too),
# so running post_available_fn would double-install.
if ok and not ran_install and post_available_fn is not None:
try:
post_available_fn(event_queue)
except Exception as exc:

9
studio/backend/tests/test_training_worker_flash_attn.py
View file

@ -692,6 +692,10 @@ def test_hook_installs_when_gate_returns_false(monkeypatch):
def test_hook_skips_install_when_gate_already_true(monkeypatch):
"""When both gates are already True AND tilelang is healthy, the hook
must do zero install work. (Tilelang repair on the already-True path
is covered by test_hook_runs_tilelang_repair_when_fla_already_true.)
"""
fla_gate = _make_fake_gate(initial_return=True)
conv_gate = _make_fake_gate(initial_return=True)
_patch_iu_gates(monkeypatch, fla_gate, conv_gate)
@ -706,6 +710,11 @@ def test_hook_skips_install_when_gate_already_true(monkeypatch):
worker, "_ensure_tilelang_backend_unconditional", tile_install
)
monkeypatch.setattr(worker, "_install_package_wheel_first", conv_install)
# Tilelang healthy so the post_available path is a no-op (otherwise
# it would call tile_install, which is correct behaviour but
# outside the scope of this test).
monkeypatch.setattr(worker, "_tilelang_importable", lambda: True)
monkeypatch.setattr(worker, "_installed_tvm_ffi_version", lambda: "0.1.9")
monkeypatch.delenv(worker._FAST_PATH_HOOKS_SKIP_ENV, raising=False)
worker._install_fast_path_hooks(event_queue=_FakeQueue(), model_name="unsloth/Qwen3.5-2B")