From 85cdafc184556a267012cd98af9f6be5fdc38f64 Mon Sep 17 00:00:00 2001 From: danielhanchen Date: Sun, 17 May 2026 01:25:25 +0000 Subject: [PATCH] studio: fix double-install of tilelang on the FLA hook install path 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. --- studio/backend/core/training/worker.py | 14 ++++++++++---- .../tests/test_training_worker_flash_attn.py | 9 +++++++++ 2 files changed, 19 insertions(+), 4 deletions(-) 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")