Studio: close switch/cancel races during model load - #6918
Conversation
Fix six race conditions when a user switches or cancels a model while a previous load or generation is still in flight, across the inference orchestrator and the /load and /unload routes: - Cancel an in-flight generation on a safetensors/MLX model switch and serialize unload with load under the inference lifecycle gate. - Cancel an in-flight load off the lifecycle gate so a Stop-loading cancel does not wait out the multi-minute load; guard the dispatched mailbox against a racing unload. - Recheck the loading marker after spawn and again after the load response before publishing, so a load cancelled mid-flight is reaped instead of going live. - Discard the loading marker before tearing the subprocess down in cancel_load, closing a spawn-after-cancel window and an orphaned compare-mode dispatcher during unload. - Match the unload target before canceling an in-flight GGUF load and add an off-gate fast path for the still-loading GGUF case. - Run the Unsloth unload off the event loop so a paused SSE stream holding _gen_lock cannot block the loop. Adds studio/backend/tests/test_orchestrator_unload_cancel.py covering the unload/cancel/switch race paths.
|
Codex usage limits have been reached for code reviews. Please check with the admins of this repo to increase the limits by adding credits. |
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request introduces a robust mechanism to cancel in-flight generations and abort loading models, preventing hangs and ensuring that outgoing models are not run to completion during unloads. It also moves model unloading off the main event loop to avoid blocking. The review feedback highlights concurrency issues with the dispatcher thread, suggesting a dedicated lock to serialize its lifecycle, and recommends using getattr or hasattr checks in the cancel_load teardown path to safely handle partially-constructed objects.
Important
The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.
| self._gen_lock = threading.Lock() # Serializes generation | ||
| # Set during a switch so a generation winning the _gen_lock handoff bails | ||
| # instead of starting on the outgoing model. | ||
| self._unload_pending = False |
There was a problem hiding this comment.
Initialize self._dispatcher_lifecycle_lock to serialize the dispatcher thread lifecycle and prevent concurrent spawns of the dispatcher thread.
| self._gen_lock = threading.Lock() # Serializes generation | |
| # Set during a switch so a generation winning the _gen_lock handoff bails | |
| # instead of starting on the outgoing model. | |
| self._unload_pending = False | |
| self._gen_lock = threading.Lock() # Serializes generation | |
| self._dispatcher_lifecycle_lock = threading.Lock() # Serializes dispatcher start/stop | |
| # Set during a switch so a generation winning the _gen_lock handoff bails | |
| # instead of starting on the outgoing model. | |
| self._unload_pending = False |
| dispatcher_preexisting = ( | ||
| self._dispatcher_thread is not None and self._dispatcher_thread.is_alive() | ||
| ) | ||
| self._start_dispatcher() |
There was a problem hiding this comment.
The _start_dispatcher method is not thread-safe and can be called concurrently by multiple threads in compare mode. If multiple threads call _generate_dispatched concurrently, they may both see self._dispatcher_thread as inactive and spawn multiple dispatcher threads. This will cause them to compete for self._resp_queue reads, leading to split responses and hung requests. We should serialize the dispatcher thread lifecycle using a dedicated lock.
| dispatcher_preexisting = ( | |
| self._dispatcher_thread is not None and self._dispatcher_thread.is_alive() | |
| ) | |
| self._start_dispatcher() | |
| with self._dispatcher_lifecycle_lock: | |
| dispatcher_preexisting = ( | |
| self._dispatcher_thread is not None and self._dispatcher_thread.is_alive() | |
| ) | |
| self._start_dispatcher() |
| if orphaned_dispatcher: | ||
| self._stop_dispatcher() |
There was a problem hiding this comment.
Acquire self._dispatcher_lifecycle_lock before calling _stop_dispatcher to ensure thread-safe serialization of the dispatcher thread lifecycle.
| if orphaned_dispatcher: | |
| self._stop_dispatcher() | |
| if orphaned_dispatcher: | |
| with self._dispatcher_lifecycle_lock: | |
| self._stop_dispatcher() |
| target = model_name | ||
| if target not in self.loading_models: | ||
| target = next( | ||
| (m for m in self.loading_models if m.lower() == model_name.lower()), | ||
| model_name, | ||
| ) | ||
| self._shutdown_subprocess(timeout = 0.5) | ||
| self.loading_models.discard(model_name) | ||
| self.active_model_name = None | ||
| self.models.clear() | ||
| if target not in self.loading_models: | ||
| return False |
There was a problem hiding this comment.
The cancel_load method is a teardown/cleanup path. To tolerate partially-constructed objects (e.g., when __init__ raises an exception before all attributes are initialized, or when the object is instantiated via __new__ in tests), we should access instance attributes using getattr instead of direct access to prevent AttributeError.
target = model_name
loading_models = getattr(self, "loading_models", set())
if target not in loading_models:
target = next(
(m for m in loading_models if m.lower() == model_name.lower()),
model_name,
)
if target not in loading_models:
return FalseReferences
- Ensure that cleanup or teardown paths tolerate partially-constructed objects by accessing instance attributes using getattr(self, 'attribute_name', None) instead of direct access to prevent AttributeError.
| self.loading_models.discard(target) | ||
| self.active_model_name = None | ||
| self.models.clear() |
There was a problem hiding this comment.
Use getattr or hasattr checks to safely discard/clear attributes on partially-constructed objects during teardown.
| self.loading_models.discard(target) | |
| self.active_model_name = None | |
| self.models.clear() | |
| if hasattr(self, "loading_models"): | |
| self.loading_models.discard(target) | |
| if hasattr(self, "active_model_name"): | |
| self.active_model_name = None | |
| if hasattr(self, "models"): | |
| self.models.clear() |
References
- Ensure that cleanup or teardown paths tolerate partially-constructed objects by accessing instance attributes using getattr(self, 'attribute_name', None) instead of direct access to prevent AttributeError.
| self.active_model_name = None | ||
| self.models.clear() |
There was a problem hiding this comment.
Use hasattr checks to safely clear attributes on partially-constructed objects during teardown.
| self.active_model_name = None | |
| self.models.clear() | |
| if hasattr(self, "active_model_name"): | |
| self.active_model_name = None | |
| if hasattr(self, "models"): | |
| self.models.clear() |
References
- Ensure that cleanup or teardown paths tolerate partially-constructed objects by accessing instance attributes using getattr(self, 'attribute_name', None) instead of direct access to prevent AttributeError.
1 participant
Summary
Closes six race conditions that surface when a user switches models or hits Stop-loading while a previous load or generation is still in flight. The fixes span the inference orchestrator and the
/loadand/unloadroutes so that a cancelled or superseded load is reliably reaped instead of going live or wedging the UI.This is the same-repo version of the fork-hosted PR #6783 (branch
studio-switch-cancelondanielhanchen/unsloth-staging-2), rebased onto currentmainso it can land here. The orchestrator and worker changes apply unchanged; the/unloadand GGUF-load route changes were re-integrated on top of the tool-calling work that has since merged intomain. It supersedes #6783.The six fixes
/unloadwith/loadunder the inference lifecycle gate so a concurrent load cannot swap in a fresh subprocess mid-unload.cancel_load, closing a spawn-after-cancel window and an orphaned compare-mode dispatcher during unload._gen_lockcannot block the loop.Tests
Adds
studio/backend/tests/test_orchestrator_unload_cancel.py(42 tests) covering the unload/cancel/switch race paths.Neighboring orchestrator suites (
test_inference_orchestrator_crash_message.py,test_openai_auto_switch.py) still pass (149 passed).