fix(studio): load run.py by path for editable installs (#5909)

* fix(studio): load run.py by path for editable installs

`studio update` can leave a partial site-packages/studio/backend/ tree
(plugin build artefacts only). That shadowed tree wins over an editable
install and breaks `from studio.backend.run import ...`. Loading run.py
by file path via importlib sidesteps the conflict.

The module is cached in _RUN_MODULE so repeated calls are cheap.
If exec_module fails, the module is removed from sys.modules before
re-raising so a subsequent retry starts clean.

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Handle None __file__ when checking cached run module for PR #5909

* Harden _load_backend_auth_storage against None __file__ and resolve cache-key path (PR #5909)

* Adapt studio run/cloudflare in-venv tests to _load_run_module loader (PR #5909)

---------

Co-authored-by: Jim Dawdy <jimdawdy@Jims-MacBook-Pro.local>
Co-authored-by: Claude Sonnet 4.6 <noreply@anthropic.com>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: Daniel Han <danielhanchen@gmail.com>

unsloth_cli/commands/studio.py

@@ -184,6 +184,47 @@ def _find_run_py() -> Optional[Path]:
     return None
 
 
+_RUN_MODULE = None
+
+
+def _load_run_module():
+    """Import studio.backend.run without relying on package resolution.
+
+    `studio update` can leave a partial ``site-packages/studio/backend/``
+    tree (plugin build artefacts only). That shadowed tree wins over an
+    editable install and breaks ``from studio.backend.run import ...``.
+    Loading by file path sidesteps the conflict.
+    """
+    global _RUN_MODULE
+    if _RUN_MODULE is not None:
+        return _RUN_MODULE
+
+    run_py = _find_run_py()
+    if run_py is None:
+        raise ImportError("Could not find studio/backend/run.py. Re-run: unsloth studio setup")
+
+    loaded = sys.modules.get("studio.backend.run")
+    if loaded is not None:
+        # __file__ can be None for namespace packages from partial trees.
+        loaded_path = Path(getattr(loaded, "__file__", None) or "").resolve()
+        if loaded_path == run_py.resolve():
+            _RUN_MODULE = loaded
+            return _RUN_MODULE
+
+    spec = importlib.util.spec_from_file_location("studio.backend.run", run_py)
+    if spec is None or spec.loader is None:
+        raise ImportError(f"Could not load studio backend from {run_py}")
+    module = importlib.util.module_from_spec(spec)
+    sys.modules["studio.backend.run"] = module
+    try:
+        spec.loader.exec_module(module)
+    except Exception:
+        sys.modules.pop("studio.backend.run", None)
+        raise
+    _RUN_MODULE = module
+    return _RUN_MODULE
+
+
 def _find_setup_script() -> Optional[Path]:
     """Find studio/setup.sh or studio/setup.ps1.
 
@@ -329,9 +370,11 @@ def _load_backend_auth_storage():
     auth_dir = backend_dir / "auth"
     storage_py = auth_dir / "storage.py"
     loaded = sys.modules.get("auth.storage")
-    loaded_path = Path(getattr(loaded, "__file__", "")).resolve()
-    if loaded is not None and loaded_path == storage_py:
-        return loaded
+    if loaded is not None:
+        # __file__ can be None for namespace packages from partial trees.
+        loaded_path = Path(getattr(loaded, "__file__", None) or "").resolve()
+        if loaded_path == storage_py.resolve():
+            return loaded
 
     package = sys.modules.get("auth")
     package_paths = [Path(path).resolve() for path in getattr(package, "__path__", [])]
@@ -706,11 +749,11 @@ def studio_default(
             typer.echo("Studio not set up. Run install.sh first.")
             raise typer.Exit(1)
 
-    from studio.backend.run import run_server
+    run_mod = _load_run_module()
+    run_server = run_mod.run_server
 
     if not silent:
-        from studio.backend.run import _resolve_external_ip
-        display_host = _resolve_external_ip() if host == "0.0.0.0" else host
+        display_host = run_mod._resolve_external_ip() if host == "0.0.0.0" else host
         typer.echo(f"Starting Unsloth Studio on http://{display_host}:{port}")
 
     run_kwargs = dict(
@@ -725,20 +768,17 @@ def studio_default(
         run_kwargs["frontend_path"] = frontend
     run_server(**run_kwargs)
 
-    from studio.backend.run import _shutdown_event
-
     try:
-        if _shutdown_event is not None:
+        if run_mod._shutdown_event is not None:
             # Event.wait() with no timeout blocks at C-level on Linux
             # and swallows SIGINT; loop with a 1s timeout instead.
-            while not _shutdown_event.is_set():
-                _shutdown_event.wait(timeout = 1)
+            while not run_mod._shutdown_event.is_set():
+                run_mod._shutdown_event.wait(timeout = 1)
         else:
             while True:
                 time.sleep(1)
     except KeyboardInterrupt:
-        from studio.backend.run import _graceful_shutdown, _server
-        _graceful_shutdown(_server)
+        run_mod._graceful_shutdown(run_mod._server)
         typer.echo("\nShutting down...")
 
 
@@ -1036,7 +1076,8 @@ def run(
             os.execvp(str(studio_bin), args)
 
     # ── 2. Start server (always suppress built-in banner) ─────────────
-    from studio.backend.run import run_server, _resolve_external_ip
+    run_mod = _load_run_module()
+    run_server = run_mod.run_server
 
     run_kwargs = dict(
         host = host,
@@ -1099,7 +1140,7 @@ def run(
     context_length_line = _format_context_length_line(result)
 
     # 6. Print banner.
-    display_host = _resolve_external_ip() if host == "0.0.0.0" else host
+    display_host = run_mod._resolve_external_ip() if host == "0.0.0.0" else host
     base_url = f"http://{display_host}:{actual_port}"
     sdk_base_url = f"{base_url}/v1"
     # run_server started the tunnel during the silent run above (0.0.0.0 only).
@@ -1178,17 +1219,15 @@ def run(
         typer.secho(_tool_notice, fg = _tool_notice_fg, bold = True)
 
     # 7. Wait for Ctrl+C.
-    from studio.backend.run import _shutdown_event, _graceful_shutdown, _server
-
     try:
-        if _shutdown_event is not None:
-            while not _shutdown_event.is_set():
-                _shutdown_event.wait(timeout = 1)
+        if run_mod._shutdown_event is not None:
+            while not run_mod._shutdown_event.is_set():
+                run_mod._shutdown_event.wait(timeout = 1)
         else:
             while True:
                 time.sleep(1)
     except KeyboardInterrupt:
-        _graceful_shutdown(_server)
+        run_mod._graceful_shutdown(run_mod._server)
         typer.echo("\nShutting down...")
 
 

unsloth_cli/tests/test_studio_cloudflare_flag.py

@@ -210,6 +210,9 @@ def test_run_in_venv_passes_cloudflare_to_run_server(monkeypatch, user_flag, exp
     )
     fake_backend_run.run_server = fake_run_server
     fake_backend_run._resolve_external_ip = lambda: "127.0.0.1"
+    # run() loads the backend via _load_run_module() (by file path); inject the
+    # mock as the cached run module so the stubbed run_server is used.
+    monkeypatch.setattr(studio_mod, "_RUN_MODULE", fake_backend_run)
 
     import typer as _typer
 
@@ -270,6 +273,9 @@ def test_run_in_venv_shuts_down_on_startup_abort(monkeypatch):
     backend._server = object()
     backend._shutdown_event = None
     backend._graceful_shutdown = lambda server: shutdown_calls.append(server)
+    # run() loads the backend via _load_run_module() (by file path); inject the
+    # mock as the cached run module so the stubbed symbols are used.
+    monkeypatch.setattr(studio_mod, "_RUN_MODULE", backend)
 
     # set_tool_policy is imported as `from state.tool_policy import set_tool_policy`.
     state_mod = sys.modules.setdefault("state", types.ModuleType("state"))

unsloth_cli/tests/test_studio_run_parallel_flag.py

@@ -426,6 +426,10 @@ def test_in_venv_path_passes_parallel_to_run_server(monkeypatch, value):
     )
     fake_backend_run.run_server = fake_run_server
     fake_backend_run._resolve_external_ip = lambda: "127.0.0.1"
+    # run() loads the backend via _load_run_module() (by file path), which
+    # ignores a sys.modules mock with no matching __file__; inject it as the
+    # cached run module so the stubbed run_server is used.
+    monkeypatch.setattr(studio_mod, "_RUN_MODULE", fake_backend_run)
 
     import typer as _typer