diff --git a/studio/backend/core/inference/orchestrator.py b/studio/backend/core/inference/orchestrator.py index d1e3c27a2a5..50f736399db 100644 --- a/studio/backend/core/inference/orchestrator.py +++ b/studio/backend/core/inference/orchestrator.py @@ -303,6 +303,10 @@ def __init__(self): # orphaning the extra thread (self._dispatcher_thread tracks only the last). The # orphan later steals the "unloaded" reply off resp_queue and hangs unload_model. self._dispatcher_lifecycle_lock = threading.Lock() + # Set under _dispatcher_lifecycle_lock while share_distributed_object reads resp_queue + # under _gen_lock with no mailbox registered: a dispatcher spawned meanwhile would take + # its reply and drop it as unaddressed. Blocks _start_dispatcher/_generate_dispatched. + self._exclusive_op_pending = False # Local state mirrors (updated from subprocess responses) self.active_model_name: Optional[str] = None @@ -1157,6 +1161,7 @@ def _start_dispatcher(self) -> bool: if ( self._unload_pending or self._exclusive_tts_pending + or self._exclusive_op_pending or getattr(self, "_exclusive_vram_probe_pending", False) ): return False @@ -1293,6 +1298,11 @@ def _generate_dispatched( yield GenStreamError("Error: audio generation is in progress", public = True) return + # A share owns resp_queue; the under-lock recheck below covers a flag set after here. + if self._exclusive_op_pending: + yield GenStreamError("Error: a distributed object share is in progress", public = True) + return + # Ensure the dispatcher runs. _start_dispatcher serializes concurrent starters under # _dispatcher_lifecycle_lock and returns True only for the caller that actually spawned # the thread, so at most one dispatcher ever exists even when two compare requests race @@ -1353,9 +1363,13 @@ def _generate_dispatched( or self.active_model_name != expected_model or not dispatcher_alive ) + # Under _mailbox_lock so the gate is atomic with registration: the share's + # _wait_dispatcher_idle stops the dispatcher, orphaning a mailbox registered in the + # window the unlocked pre-check leaves. Separate from `unloading` so the refusal names it. + share_reserved = self._exclusive_op_pending tts_reserved = self._exclusive_tts_pending probe_reserved = getattr(self, "_exclusive_vram_probe_pending", False) - blocked = unloading or tts_reserved or probe_reserved + blocked = unloading or share_reserved or tts_reserved or probe_reserved if not blocked: self._mailboxes[request_id] = mailbox if cancel_event is not None: @@ -1373,7 +1387,9 @@ def _generate_dispatched( if orphaned_dispatcher: self._stop_dispatcher() detail = ( - "Error: audio generation is in progress" + "Error: a distributed object share is in progress" + if share_reserved + else "Error: audio generation is in progress" if tts_reserved else "Error: model switch is checking GPU memory" if probe_reserved @@ -1492,49 +1508,61 @@ def share_distributed_object( if not self._ensure_subprocess_alive(): raise RuntimeError("Inference subprocess is not running") - self._wait_dispatcher_idle() - with self._mailbox_lock: - if self._mailboxes: + # Set BEFORE the drain: the window between it and the read loop below is spawnable. + with self._dispatcher_lifecycle_lock: + self._exclusive_op_pending = True + try: + # False means it left the dispatcher alive, and only this return says so: a compare + # stream unregistering just past the deadline empties the mailbox snapshot below. + if not self._wait_dispatcher_idle(): raise RuntimeError( "Cannot share distributed objects while compare requests are active" ) - request_id = str(uuid.uuid4()) - cmd = { - "type": "share_object", - "request_id": request_id, - "object": obj, - } + with self._mailbox_lock: + if self._mailboxes: + raise RuntimeError( + "Cannot share distributed objects while compare requests are active" + ) + request_id = str(uuid.uuid4()) + cmd = { + "type": "share_object", + "request_id": request_id, + "object": obj, + } - with self._gen_lock: - self._send_cmd(cmd) - deadline = None if timeout is None else time.monotonic() + timeout - while deadline is None or time.monotonic() < deadline: - remaining = 1.0 if deadline is None else max(0.1, deadline - time.monotonic()) - resp = self._read_resp(timeout = min(remaining, 1.0)) - if resp is None: - if not self._ensure_subprocess_alive(): - raise RuntimeError(self._subprocess_crash_message("sharing chat turn")) - continue + with self._gen_lock: + self._send_cmd(cmd) + deadline = None if timeout is None else time.monotonic() + timeout + while deadline is None or time.monotonic() < deadline: + remaining = 1.0 if deadline is None else max(0.1, deadline - time.monotonic()) + resp = self._read_resp(timeout = min(remaining, 1.0)) + if resp is None: + if not self._ensure_subprocess_alive(): + raise RuntimeError(self._subprocess_crash_message("sharing chat turn")) + continue - rtype = resp.get("type", "") - rid = resp.get("request_id") - if rid and rid != request_id: - logger.debug( - "Skipping response for request_id=%s while sharing request_id=%s", - rid, - request_id, - ) - continue - if rtype == "shared": - return resp.get("object") - if rtype == "share_error": - raise RuntimeError(resp.get("error", "Failed to share object")) - if rtype == "error": - raise RuntimeError(resp.get("error", "Subprocess error")) - if rtype == "status": - continue + rtype = resp.get("type", "") + rid = resp.get("request_id") + if rid and rid != request_id: + logger.debug( + "Skipping response for request_id=%s while sharing request_id=%s", + rid, + request_id, + ) + continue + if rtype == "shared": + return resp.get("object") + if rtype == "share_error": + raise RuntimeError(resp.get("error", "Failed to share object")) + if rtype == "error": + raise RuntimeError(resp.get("error", "Subprocess error")) + if rtype == "status": + continue - raise RuntimeError("Timeout waiting for distributed object share") + raise RuntimeError("Timeout waiting for distributed object share") + finally: + with self._dispatcher_lifecycle_lock: + self._exclusive_op_pending = False # ------------------------------------------------------------------ # Public API — same interface as InferenceBackend diff --git a/studio/backend/routes/settings.py b/studio/backend/routes/settings.py index a3197ffa14a..c2d0a86c76e 100644 --- a/studio/backend/routes/settings.py +++ b/studio/backend/routes/settings.py @@ -2748,12 +2748,23 @@ def update_embedding_model( ) ) ) - if evaluate_file_security( - verify_target, - hf_token = scan_token, - load_subdirs = load_subdirs, - local_only_load = local_only_load, - ).blocked: + try: + _security_blocked = evaluate_file_security( + verify_target, + hf_token = scan_token, + load_subdirs = load_subdirs, + local_only_load = local_only_load, + ).blocked + except Exception: + # A scan error is a gate failure, not a verdict: fail open rather than 500 the + # route, as _guard_model_security in core/rag/embeddings.py does. + logger.warning( + "Embedding-model security scan errored for %r; allowing (fail-open)", + model, + exc_info = True, + ) + _security_blocked = False + if _security_blocked: # 403, not 409: the client routes every 409 into the forceable "save anyway" # flow, but this block is a hard, non-forceable security refusal. if local_only_load: diff --git a/studio/backend/tests/test_audio_tts_cancellation.py b/studio/backend/tests/test_audio_tts_cancellation.py index ef19f1ec313..1b329a7ff37 100644 --- a/studio/backend/tests/test_audio_tts_cancellation.py +++ b/studio/backend/tests/test_audio_tts_cancellation.py @@ -92,6 +92,7 @@ def _bare_orchestrator(): orchestrator._request_cancel_events = {} orchestrator._unload_pending = False orchestrator._exclusive_tts_pending = False + orchestrator._exclusive_op_pending = False orchestrator.active_model_name = "model" orchestrator.models = {"model": {}} orchestrator.loading_models = set() diff --git a/studio/backend/tests/test_embedding_model_security_gate.py b/studio/backend/tests/test_embedding_model_security_gate.py index 1bb5ca09c36..ff7df7aafa4 100644 --- a/studio/backend/tests/test_embedding_model_security_gate.py +++ b/studio/backend/tests/test_embedding_model_security_gate.py @@ -115,6 +115,31 @@ def test_flagged_repo_is_blocked_without_force(client, monkeypatch): assert "model" not in saved +def _security_raises(): + mod = _types.ModuleType("utils.security") + + def _boom(*_a, **_k): + raise RuntimeError("scan endpoint unreachable") + + mod.evaluate_file_security = _boom + mod.security_load_subdirs = lambda *a, **k: () + return mod + + +def test_scan_error_fails_open_instead_of_500(client, monkeypatch): + # A scan error is a gate failure, not a verdict; same policy as _guard_model_security. + c, saved = client + monkeypatch.setitem(sys.modules, "utils.security", _security_raises()) + import utils.models as models + + monkeypatch.setattr(models, "is_embedding_model", lambda *a, **k: True) + + r = c.put("/embedding-model", json = {"embedding_model": "acme/embedder"}) + + assert r.status_code == 200 + assert saved["model"] == "acme/embedder" + + def test_uncached_selection_is_marked_pending_so_loaders_stay_offline(client, monkeypatch): c, saved = client monkeypatch.setitem(sys.modules, "utils.security", _security_stub(blocked = False)) diff --git a/studio/backend/tests/test_mcp_session_resources.py b/studio/backend/tests/test_mcp_session_resources.py index bae4252dd0c..71a192f995f 100644 --- a/studio/backend/tests/test_mcp_session_resources.py +++ b/studio/backend/tests/test_mcp_session_resources.py @@ -133,9 +133,9 @@ def test_repeated_open_close_does_not_accumulate_threads(tiny): for i in range(12): call_tool_sync(HTTP_URL, None, "t", {}, scope = f"chat-{i}") close_mcp_sessions() - assert _settle(lambda: _session_threads() <= before), ( - f"leaked threads after 12 cycles: {_session_threads()} vs {before}" - ) + assert _settle( + lambda: _session_threads() <= before + ), f"leaked threads after 12 cycles: {_session_threads()} vs {before}" def test_repeated_open_close_does_not_accumulate_descriptors(tiny): diff --git a/studio/backend/tests/test_mcp_upgrade_compat.py b/studio/backend/tests/test_mcp_upgrade_compat.py index 4ef03ed9fe3..434bc698be0 100644 --- a/studio/backend/tests/test_mcp_upgrade_compat.py +++ b/studio/backend/tests/test_mcp_upgrade_compat.py @@ -107,9 +107,9 @@ def test_only_stdio_answers_the_liveness_probe(): from fastmcp.client.transports import SSETransport, StreamableHttpTransport for cls in (StreamableHttpTransport, SSETransport): transport = cls(url = "https://x.test/mcp") - assert not hasattr(transport, "_is_session_dead"), ( - f"{cls.__name__} grew a liveness probe; _transport_dead can use it now" - ) + assert not hasattr( + transport, "_is_session_dead" + ), f"{cls.__name__} grew a liveness probe; _transport_dead can use it now" def test_the_installed_fastmcp_meets_the_declared_floor(): diff --git a/studio/backend/tests/test_orchestrator_unload_cancel.py b/studio/backend/tests/test_orchestrator_unload_cancel.py index 966b5e06a94..d73ec224154 100644 --- a/studio/backend/tests/test_orchestrator_unload_cancel.py +++ b/studio/backend/tests/test_orchestrator_unload_cancel.py @@ -34,6 +34,7 @@ def _bare_orchestrator(): o._dispatcher_lifecycle_lock = threading.Lock() o._unload_pending = False o._exclusive_tts_pending = False + o._exclusive_op_pending = False o.active_model_name = "m" o.models = {"m": {}} o.loading_models = set() @@ -2646,3 +2647,64 @@ async def cancel_then_die(*args, **kwargs): with inf._scoped_load_attempts_lock: inf._scoped_load_attempts.clear() inf._scoped_load_cancel_tombstones.clear() + + +def test_share_aborts_when_dispatcher_drain_fails(monkeypatch): + # Reading _mailboxes instead: a stream unregistering just past the deadline empties the + # map, so the share runs on while that dispatcher eats its reply. timeout=None hangs. + o = _bare_orchestrator() + o._mailbox_lock = threading.Lock() + o._mailboxes = {} + monkeypatch.setattr(o, "_ensure_subprocess_alive", lambda: True) + monkeypatch.setattr(o, "_wait_dispatcher_idle", lambda: False) + monkeypatch.setattr( + o, "_send_cmd", lambda cmd: pytest.fail("must not share past a failed drain") + ) + + with pytest.raises(RuntimeError, match = "compare requests are active"): + o.share_distributed_object({"role": "user"}, timeout = 1.0) + + assert o._exclusive_op_pending is False, "the exclusive flag must not leak on the abort" + + +def test_dispatched_share_refusal_is_public(monkeypatch): + # Without public=True _friendly_gen_stream_error swaps this for "An internal error occurred." + o = _bare_orchestrator() + monkeypatch.setattr(o, "_ensure_subprocess_alive", lambda: True) + monkeypatch.setattr( + o, "_start_dispatcher", lambda: pytest.fail("must not start a dispatcher during a share") + ) + o._exclusive_op_pending = True + + out = list(o._generate_dispatched(messages = [{"role": "user", "content": "hi"}])) + + assert len(out) == 1 + assert "distributed object share" in str(out[0]) + assert out[0].public is True + + +def test_dispatched_share_recheck_refusal_names_the_share(monkeypatch): + # Folding this into `unloading` said "model is being unloaded" with no unload in sight. + o = _bare_orchestrator() + o._mailbox_lock = threading.Lock() + o._mailboxes = {} + o._request_cancel_events = {} + o._dispatcher_thread = _AliveDispatcher() + monkeypatch.setattr(o, "_ensure_subprocess_alive", lambda: True) + monkeypatch.setattr(o, "_start_dispatcher", lambda: False) + monkeypatch.setattr( + o, "_send_cmd", lambda cmd: pytest.fail("must not send generate during a share") + ) + + def flip(*a, **k): + o._exclusive_op_pending = True + return {"type": "generate", "request_id": "r1"} + + monkeypatch.setattr(o, "_build_generate_cmd", flip) + + out = list(o._generate_dispatched(messages = [{"role": "user", "content": "hi"}])) + + assert len(out) == 1 + assert "distributed object share" in str(out[0]) + assert "unloaded" not in str(out[0]) + assert o._mailboxes == {}, "must not leave an orphaned mailbox"