diff --git a/studio/backend/core/training/worker.py b/studio/backend/core/training/worker.py index 9e4fe37b5e..3b48856eb7 100644 --- a/studio/backend/core/training/worker.py +++ b/studio/backend/core/training/worker.py @@ -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: diff --git a/studio/backend/tests/test_training_worker_flash_attn.py b/studio/backend/tests/test_training_worker_flash_attn.py index 2cfe512be2..78cad64892 100644 --- a/studio/backend/tests/test_training_worker_flash_attn.py +++ b/studio/backend/tests/test_training_worker_flash_attn.py @@ -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")